/* stitch.c : stitches files together */
/* TODO:
 * + support a CRC option
 *
 * BUGS:
 * + -i options cannot overlap
 * + -i option does not work if it crosses a 16K boundry.
 */

#include <assert.h>
#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/***** Debug/tracing/error macros *****/

#define DIE_ON(x, ...) do { if(x) { fprintf(stderr, __VA_ARGS__); exit(EXIT_FAILURE); } } while(0)

#ifndef NDEBUG
#define PERROR(reason) fprintf(stderr, "%s:%d:%s():%s:%s\n", __FILE__, __LINE__, __func__, reason, strerror(errno))
#define DIE_ON_ERROR(x, reason) DIE_ON((x), "%s:%s:" #x "\n", reason, strerror(errno))
#else
#define PERROR(reason) fprintf(stderr, "%s:%s\n", reason, strerror(errno))
#define DIE_ON_ERROR(x, reason) DIE_ON((x), ":%s:%s\n", reason, strerror(errno))
#endif

/***** TYPES *****/

struct opt {
	unsigned pad_size; /* 0 = no padding */
	unsigned char fill_char;
	int endian; /* 0 == big endian, 1 = little endian */
};

struct file_patch {
	char *filename;
	FILE *f;
};

struct integer_entry {
	unsigned long offset;
	unsigned value;
	int endian;
	unsigned size;
	struct integer_entry *next;
};

struct patch_entry {
	int type; /* 0 = file, 1 = empty */
	unsigned long offset;
	unsigned length; /* in bytes */
	struct opt options;
	struct file_patch file;
	struct patch_entry *next;
};

/***** GLOBALS *****/

static const char *program_name;
static const char *outfile;
static struct patch_entry *patch_list_head;
static struct integer_entry *integer_list_head;
static unsigned long tail_ofs; /* track the end */

/***** FUNCTIONS : opt *****/

static void opt_default(struct opt *op) {
	op->pad_size=0;
	op->fill_char=0xff;
}

static unsigned long long suffix_multiplier(const char *suffix) {
	assert(suffix != NULL);
	if(suffix[0]==0) {
		return 1;
	}
	if(suffix[1]==0 || suffix[1]=='b' || suffix[1]=='B' || suffix[1]=='i' || suffix[1]=='I') {
		switch(suffix[0]) {
			case 'k': case 'K': return 0x400;
			case 'm': case 'M': return 0x100000;
			case 'g': case 'G': return 0x40000000;
			case 't': case 'T': return 0x10000000000ll;
			case 'p': case 'P': return 0x4000000000000ll;
			case 'e': case 'E': return 0x1000000000000000ll;
			/*
			case 'z': case 'Z': return 0x400000000000000000;
			case 'y': case 'Y': return 0x100000000000000000000;
			*/
		}
	}
	DIE_ON(1, "Invalid suffix '%s'\n", suffix);
}

static void opt_padsize(struct opt *op, const char *s) {
	char *ep;
	unsigned long n;
	n=strtoul(s, &ep, 0); /* decimal or hex */
	DIE_ON(!*s || s==ep, "pad size does not make sense.\n");
	op->pad_size=n*suffix_multiplier(ep);
}

static void opt_endian(struct opt *op, const char *s) {
	if(toupper(*s)=='B' || *s=='0') {
		op->endian=0; /* Big Endian */
	} else if(toupper(*s)=='L' || *s=='1') {
		op->endian=1; /* Little Endian */
	} else {
		DIE_ON(1, "Endian option(-E) does not make sense.\n");
	}
}

static void opt_fillchar(struct opt *op, const char *s) {
	char *ep;
	unsigned long n;
	if(s[0]!=0 && s[1]==0) {
		n=s[0];
	} else {
		s++;
		n=strtoul(s, &ep, 0); /* decimal or hex */
		DIE_ON(*ep, "fill character does not make sense");
	}

	op->fill_char=n;
}

/***** FUNCTIONS : patch_list *****/
static struct patch_entry *patch_listadd(unsigned type, long offset, unsigned long length, struct opt *options) {
	struct patch_entry *pe, *curr, **prev;
	if(offset<0) {
		offset=tail_ofs;
	}
	if(offset+length>tail_ofs) {
		tail_ofs+=offset+length;
	}

	pe=malloc(sizeof *pe);
	pe->type=type;
	pe->offset=offset;
	pe->length=length;
	if(options) {
		pe->options=*options;
	} else {
		opt_default(&pe->options);
	}

	for(prev=&patch_list_head,curr=*prev;curr;prev=&curr->next,curr=*prev) {
		/* TODO: find a position to insert the new entry */
		if(curr->offset>=(unsigned)offset) break;
	}
	*prev=pe;
	pe->next=curr;
	return pe;
}

static void patch_list_free(void) {
	struct patch_entry *curr, *tmp;

	for(curr=patch_list_head;curr;curr=tmp) {
		tmp=curr->next;
		if(curr->type==0) {
			fclose(curr->file.f);
			free(curr->file.filename);
		}
		free(curr);
	}
	patch_list_head=0;
}

/***** FUNCTIONS : integer_entry *****/

static int parse_var_value(const char *str, unsigned long *address, unsigned long *value) {
	char *endptr;
	unsigned long addr, val;
	if(!str || !*str) return 0; /* failure */
	addr=strtoul(str, &endptr, 16); /* address is always hex */
	if(endptr==str) return 0; /* failure */
	str=endptr;
	if(*str!='=') return 0; /* failure */
	val=strtoul(str+1, &endptr, 0); /* value autodetects */
	if(*endptr) return 0; /* failure */
	if(address) *address=addr;
	if(value) *value=val;
	return 1; /* success */
}

static struct integer_entry *integer_listadd(unsigned long address, unsigned long value, int endian) {
	struct integer_entry *ie, *curr, **prev;

	ie=malloc(sizeof *ie);
	ie->offset=address;
	ie->value=value;
	ie->endian=endian;
	ie->size=4;
	for(prev=&integer_list_head,curr=*prev;curr;prev=&curr->next,curr=*prev) {
		/* TODO: find a position to insert the new entry */
		if(curr->offset>=(unsigned)address) break;
	}
	*prev=ie;
	ie->next=curr;
	return ie;
}

static void integer_s_listadd(const char *str, int endian) {
	unsigned long address, value;
	if(parse_var_value(str, &address, &value)) 
		integer_listadd(address, value, endian);
	else {
		fprintf(stderr, "ERROR: could not parse '%s'\n", str);
		exit(EXIT_FAILURE);
	}
}

/***** FUNCTIONS : empty_patch *****/

static void emptypad_listadd(const char *length, struct opt *options) {
	char *ep;
	unsigned long len;
	struct patch_entry *pe;

	len=strtoul(length, &ep, 0); /* decimal or hex */
	DIE_ON(!*length || length==ep, "pad size does not make sense.\n");
	len*=suffix_multiplier(ep);

	pe=patch_listadd(1, -1, len, options); /* 32-bits */
}

/***** FUNCTIONS : file_patch *****/

static long file_size(FILE *f, const char *filename) {
	long ret;
	fpos_t orig;
	fgetpos(f, &orig);
	fseek(f, 0, SEEK_END);
	ret=ftell(f);
	if(ret<0) {
		perror(filename);
		return ret;
	}
	fsetpos(f, &orig);
	return ret;
}

static void file_listadd(const char *filename, struct opt *options) {
	FILE *f;
	struct patch_entry *pe;
	struct opt opttmp;
	unsigned long len;

	/* fall back to default options */
	if(!options) {
		opt_default(&opttmp);
		options=&opttmp;
	}

	DIE_ON_ERROR(!(f=fopen(filename, "rb")), filename);

	len=file_size(f, filename);
	if(options->pad_size>0) {
		if(len>options->pad_size) {
			fprintf(stderr, "ERROR:file %s exceeds pad size %u\n", filename, options->pad_size);
			exit(EXIT_FAILURE);
		}
		len=options->pad_size;
	}

	pe=patch_listadd(0, -1, len, options);

	pe->file.filename=strdup(filename);
	pe->file.f=f;
}

/***** FUNCTIONS : stitching *****/

static int write_integer(struct integer_entry *ie, FILE *f, unsigned long *total_written, const char *filename) {
	char buf[8];
	unsigned i;
	int r;
	unsigned long value;

	value=ie->value;

	if(ie->endian==0) {
		/* Big Endian */
		for(i=ie->size;i>0;) {
			buf[--i]=(unsigned char)value;
			value>>=8;
		}
	} else {
		/* Little Endian */
		for(i=0;i<ie->size;i++) {
			buf[i]=(unsigned char)value;
			value>>=8;
		}
	}
	r=fwrite(buf, 1, ie->size, f);
	DIE_ON_ERROR(r<0 || ferror(f), filename);
	*total_written+=r;
	printf("(wrote %08x @ 0x%08lx)\n", ie->value, ie->offset); fflush(stdout);
	return r;
}

static int write_data(void *data, size_t len, FILE *f, unsigned long *total_written, const char *filename) {
	int count, r, sublen;

	count=0;

	while(len>0) {
		if(integer_list_head && (*total_written<=integer_list_head->offset
			&& (*total_written+len)>integer_list_head->offset)) {
			/* there is an integer entry between our range */
			sublen=integer_list_head->offset-*total_written;
			DIE_ON(sublen<0, "write_data() calculated a negative length.\n");
			if(sublen==0) { /* write out the next integer */
				struct integer_entry *ie;
				
				/* pop the integer_entry from the sorted list */
				ie=integer_list_head; 
				integer_list_head=integer_list_head->next;
				ie->next=0;

				r=write_integer(ie, f, total_written, filename);
				count+=r;
				free(ie); /* done with the integer_entry */
				/* TODO: fix this */
				DIE_ON(len<r, "TODO: we do not support -i option that crosses a 16K boundry.\n");
				len-=r;
				data=(char*)data+r;
				continue; /* back to the top of the loop */
			}
		} else {
			sublen=len; /* no patches to apply */
		}
		r=fwrite(data, 1, sublen, f);
		DIE_ON_ERROR(r<0 || ferror(f), filename);
		DIE_ON_ERROR(r==0, "fwrite returned 0 and it shouldn't have.\n");
		*total_written+=r;
		count+=r;

		/* move through the bytes */
		len-=r;
		data=(char*)data+r;
	}
	return count;
}

static void write_patch_entry(FILE *out, const char *outname, struct patch_entry *pe, unsigned long *total_written) {
	static unsigned char buf[16384]; /* temporary buffer */
	size_t len;
	int res;
	unsigned long bytes;

	switch(pe->type) {
		case 0:
			printf("%s ..", pe->file.filename);
			fflush(stdout);

			bytes=0;
			while((len=fread(buf, 1, sizeof buf, pe->file.f))>0) {
				res=write_data(buf, len, out, total_written, outname);
				bytes+=res;
			}
			DIE_ON_ERROR(ferror(pe->file.f), pe->file.filename);

			printf(" %lu bytes", bytes); fflush(stdout);

			if(pe->options.pad_size) {
				DIE_ON(bytes>pe->options.pad_size, "image %s exceeds padding\n", pe->file.filename);
				memset(buf, pe->options.fill_char, sizeof buf);
				while(bytes<pe->options.pad_size) {
					len=sizeof buf;
					if(bytes+len > pe->options.pad_size) {
						len=pe->options.pad_size-bytes;
					}
					/* fprintf(stderr, "writing %lu\n", len); */
					res=write_data(buf, len, out, total_written, outname);
					/* printf("res=%d\n", res); */
					bytes+=res;
				}
				printf(", pad %u (fill=0x%02X)", pe->options.pad_size, pe->options.fill_char);
			}

			printf("\n");
			break;
		case 1:
			if(pe->options.pad_size) {
				fprintf(stderr, "WARNING: ignoring pad size(-s) option for empty pad item(-e)\n");
			}

			printf("empty region .. %lu\n", pe->length);

			memset(buf, pe->options.fill_char, sizeof buf);
			for(bytes=0;bytes<pe->length;bytes+=res) {
				len=pe->length-bytes;
				if(len>sizeof buf) len=sizeof buf;
				res=write_data(buf, len, out, total_written, outname);
			}
			break;
		default:
			fprintf(stderr, "ERROR:Unknown patch type %u\n", pe->type);
			exit(EXIT_FAILURE);
	}
}

static void stitch(const char *filename) {
	FILE *f;
	struct patch_entry *curr;
	unsigned long total_written;

	DIE_ON_ERROR(!(f=fopen(filename, "wb")), filename);
	
	total_written=0; /* tracks the current offset in the file */
	for(curr=patch_list_head;curr;curr=curr->next) {
		write_patch_entry(f, filename, curr, &total_written);
	}

	fclose(f);
}

/***** FUNCTIONS : command parsing *****/

static void usage(int failed_fl) {
	if(failed_fl) fprintf(stderr, "Error processing parameters!\n");
	fprintf(stderr, 
		"%s [-o <filename>] [-f <0x4d | 77 | M>] [-s <padsize[KMGT]>] [-e <size[KMGT]>] [-i addr=value] [-E B | L ] [input file]\n"
		"  -o    output filename. default=out.bin\n"
		"  -f    fill char. hex, decimal or single character\n"
		"  -s    round up file size. hex or decimal with suffix Kb,Mb,Gb,Tb,..\n"
		"  -e    pad with an empty region. hex or decimal with suffix Kb,Mb,Gb,Tb,..\n"
		"  -i    integer: set 32-bit substitute value at a global address\n"
		"  -E    endian. either B or L for big or little endian.\n"
		"\n"
		"* The -i option creat substite values to use while writing.\n"
		"\nBUGS:\n"
 		"* the -i options cannot overlap.\n"
		"\n"
	, program_name);	
	if(failed_fl) exit(EXIT_FAILURE);
}

void needs_arg(int argi, int nr, int argc, char **argv) {
	if(argi+nr>=argc) {
		fprintf(stderr, "ERROR: Not enough arguments at %s\n", argv[argi]);
		usage(1);
	}
}

static void parse_args(int argc, char **argv) {
	int argi;
	struct opt options;

	opt_default(&options);
	for(argi=1;argi<argc;argi++) {
		if(argv[argi][0]=='-') {
			switch(argv[argi][1]) {
				case 'f': /* fill character */
					needs_arg(argi, 1, argc, argv);
					opt_fillchar(&options, argv[++argi]);
					break;
				case 's': /* file padding */
					needs_arg(argi, 1, argc, argv);
					opt_padsize(&options, argv[++argi]);
					break;
				case 'o': /* output file */
					needs_arg(argi, 1, argc, argv);
					DIE_ON(outfile!=NULL, "outfile file already specified\n");
					outfile=argv[++argi];
					break;
				case 'i': /* integer - big endian */
					needs_arg(argi, 1, argc, argv);
					integer_s_listadd(argv[++argi], 0);
					break;
				case 'e': /* pad with empty region */
					needs_arg(argi, 1, argc, argv);
					emptypad_listadd(argv[++argi], &options);
					opt_default(&options); /* clear the options each time */
					break;	
				case 'E': /* endian */
					needs_arg(argi, 1, argc, argv);
					opt_endian(&options, argv[++argi]);	
				default:
					usage(1);
			}
		} else {
			file_listadd(argv[argi], &options);
			opt_default(&options); /* clear the options each time */
		}
	}
}

/***** MAIN *****/
int main(int argc, char **argv) {
	program_name=argv[0];
	if(argc==1) usage(1);
	parse_args(argc, argv);
	stitch(outfile ? outfile : "out.bin");
	patch_list_free();
	return EXIT_SUCCESS;
}