#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* 
 * hex_merge.c
 *
 * Merge two or more hex files into one big hex file
 *
 * Usefull for using bootloaders with AVR, and probably other uses.
 *
 * (c) 2011 Thomas Ilnseher (illth@gmx.de)
 * Licensed under the GPLv2.
 *
 * This comes with absolutely no warranty of any kind
 */

/* use at least a 2byte type here */
typedef unsigned int data_t;

static data_t* data_buf = NULL;
static int db_len = 0;
static int real_len = 0;

static inline unsigned int hex4(char c) {
	if (c >= '0' && c <= '9')
		return c - '0';
	if (c >= 'a' && c <= 'f')
		return c - 'a' + 10;
	return (unsigned int)c - 'A' + 10;
}
static inline unsigned int get_hex8(char* line, int idx) {
	idx = idx * 2 + 1;

	return (hex4(line[idx]) << 4) | hex4(line[idx + 1]);
}

static inline int is_hex(char c) {
	if (c >= '0' && c <= '9')
		return 1;
	if (c >= 'a' && c <= 'f')
		return 1;
	if (c >= 'A' && c <= 'F')
		return 1;
	return 0;
}

void load_file(const char* file_name) {
	char line_buf[256];
	int seg_offs = 0;

	int line_len;
	int total_hex_digits;
	int i;
	int checksum;
	int record_type;
	int address;
	FILE* fp;
	int over_write_warn_on = 0;
	int warn_start_addr = 0xdead;

	fp = fopen(file_name, "r");
	if (! fp) {
		printf("Error: Failed to open %s for reading\n", file_name);
		exit(2);
	}

	line_buf[255] = '\0';

	while (!feof(fp)) {
		if (fgets(line_buf, 255, fp) == NULL)
			break;

		/* Sanity check: first char must be a ':' */
		if (line_buf[0] != ':') {
			printf("Warning: expected ':' as first char but got \"%s\" (File %s)\n", line_buf, file_name);
			continue;
		}
		/* sanity check: enough hex records, checksum, only hex chars */
		line_len = get_hex8(line_buf, 0);
		total_hex_digits = 2 * (line_len + 5);
		for (i = 1; i <= total_hex_digits; i++) {
			if (! is_hex(line_buf[i])){
				printf("Warning: non hexadezimal character at %d in line \"%s\" (File %s)\n", i, line_buf, file_name);
				break;
			}
		}
		if (i <= total_hex_digits)
			continue; /* hit the break in the loop above */

		/* checksum check */
		checksum = 0;
		for (i = 0; i < line_len + 5; i++)
			checksum += get_hex8(line_buf, i);

		if (checksum & 0xff) {
			printf("Warning: Checksum error in line \"%s\" (File %s)\n", line_buf, file_name);
			continue;
		}

		/* check for invalid record type */
		record_type = get_hex8(line_buf, 3);

		if (record_type > 5) {
			printf("Warning: Invalid record type in line \"%s\" (File %s)\n", line_buf, file_name);
			continue;
		}

		/* don't give a shit about 03 and 05 */
		if (record_type == 3 || record_type == 5)
			continue;

		if (record_type == 1) {
			if (over_write_warn_on)
				printf("Warning: overwriting data (range 0x%04x - 0x%04x) while loading file %s\n",
						warn_start_addr, real_len - 1, file_name);

			fclose(fp);
			return; /* EOF record */
		}

		if (record_type == 2) {
			/* segment address */
			if (line_len != 2) {
				printf("Warning: Invalid 02 record length in line \"%s\" (File %s)\n", line_buf, file_name);
				if (line_len < 2)
					continue;
			}
			seg_offs = (get_hex8(line_buf, 4) << 12) | ((get_hex8(line_buf, 5) & 0xf0) << 4);
			continue;
		}
		if (record_type == 4) {
			/* segment address */
			if (line_len != 2) {
				printf("Warning: Invalid 04 record length in line \"%s\" (File %s)\n", line_buf, file_name);
				if (line_len < 2)
					continue;
			}
			seg_offs = (get_hex8(line_buf, 4) << 24) | (get_hex8(line_buf, 5) << 16);
			continue;
		}

		/* so, here we have a well formated data record with correct checksum */

		address = seg_offs + (get_hex8(line_buf, 1) << 8) | get_hex8(line_buf, 2);
		while (address + line_len >= db_len) {
			data_t* new_buf = malloc(sizeof(data_t) * (db_len + 0x1000));
			if (! new_buf) {
				printf("Error: Malloc failed for 0x%x bytes!\n", (int)sizeof(data_t) * (db_len + 0x1000));
				exit(1);
			}
			if (data_buf) {
				memcpy(new_buf, data_buf, db_len * sizeof(data_t));
				free(data_buf);
			}
			data_buf = new_buf;
			memset(data_buf + db_len, 0xff, 0x1000 * sizeof(data_t));
			db_len += 0x1000;
		}
		real_len = address + line_len;

		for (i = 0; i < line_len; i++) {
			if ((data_buf[address + i] & 0xff00) == 0 && over_write_warn_on == 0) {
				over_write_warn_on = 1;
				warn_start_addr = address + i;
			}

			if ((data_buf[address + i] & 0xff00) != 0 && over_write_warn_on != 0) {
				over_write_warn_on = 0;
				printf("Warning: overwriting data (range 0x%04x - 0x%04x) while loading file %s\n",
						warn_start_addr, address + i, file_name);
			}

			data_buf[address + i] = get_hex8(line_buf, i + 4);
		}
	}
	fclose(fp);
	if (over_write_warn_on)
		printf("Warning: overwriting data (range 0x%04x - 0x%04x) while loading file %s\n",
				warn_start_addr, real_len - 1, file_name);
	printf("Warning: Leave parse loop due to no EOF record in file %s\n", file_name);
}

void write_file(FILE* fp, int strip_ff) {
	int i;
	int lba_flag = 0;
	int seg_offs = 0;
	int addr_out;
	int output_ptr = 0;
	int checksum;
	int segment;
	int record_len;

	if (real_len > 0xfffff)
		lba_flag = 1;

	while (output_ptr < real_len) {
		if (strip_ff) {
			/* the check: if this position was filled with the memset,
			 * then all bytes are filled, so this is only possible
			 * if this position was later overwritten by get_hex8
			 */
			while ((data_buf[output_ptr] & 0xff00) != 0)
				output_ptr++;
		}
		addr_out = output_ptr - seg_offs;
		if (addr_out > 0xfff0) {
			/* need to adjust the segment */
			if (lba_flag) {
				/* wrtie type 4 record */
				segment = ((unsigned int)output_ptr >> 16) & 0xffff;
				checksum = - (2 + 0 + 4 + (segment >> 8) + segment);
				fprintf(fp, ":02000004%04X%02X\r\n", segment, checksum & 0xff);
				seg_offs = segment << 16;
			} else {
				/* write type 2 record */
				segment = (((unsigned int)output_ptr >> 16) & 0xf) << 12;
				checksum = - (2 + 0 + 2 + (segment >> 8) + segment);
				fprintf(fp, ":02000002%04X%02X\r\n", segment, checksum & 0xff);
				seg_offs = segment << 4;
			}
		}
		addr_out = output_ptr - seg_offs;
		for (record_len = 1; record_len < 16 && ((data_buf[output_ptr + record_len] & 0xff00) == 0); record_len++)
			;

		if (real_len - output_ptr < record_len) 
			record_len = real_len - output_ptr;


		fprintf(fp, ":%02X%04X00", record_len, addr_out);
		checksum = -(record_len + (addr_out >> 8) + addr_out + 0);
		for (i = 0; i < record_len; i++) {
			fprintf(fp, "%02X", data_buf[output_ptr] & 0xff);
			checksum -= (unsigned int)(data_buf[output_ptr]);
			output_ptr++;
		}
		fprintf(fp, "%02X\r\n", checksum & 0xff);
	}

	/* print EOF record */
	fprintf(fp, ":00000001FF\r\n");
}

void usage(const char* prog) {
	printf("USAGE:\n\n%s [options] file1 [...] [output file]\n\nOptions:\n\t -o\t--output=FILE\tSte the output file (instead of the last file named)\n", prog);
	printf("\t -f\t--fill\t\tFill the empty spaces (default: off)\n");
	printf("\t -F\t--fillbyte=0xXX\tFill the empty spaces with byte 0xXX (also enables fill)\n");
	printf("\t\t\t\tDefault fill byte: 0xFF\n");
	printf("\n\nEXAMPLES:\n\n");
	printf("\t%s file1.hex file2.hex output.hex\t\t\tMerge file1.hex and file2.hex into output.hex (most basic example)\n", prog);
	printf("\t%s file1.hex file2.hex file3.hex output.hex\t\tMerge file1.hex, file2.hex and file3.hex into output.hex\n", prog);
	printf("\t%s --output=output.hex file1.hex file2.hex\t\tMerge file1.hex and file2.hex into output.hex\n", prog);
	printf("\t%s -oF output.hex 0x12 file1.hex file2.hex\t\tMerge file1.hex and file2.hex into output.hex, filling empty spaces with 0x12\n", prog);
	printf("\t%s --fillbyte=0x12 file1.hex file2.hex output.hex\tMerge file1.hex and file2.hex into output.hex, filling empty spaces with 0x12\n", prog);
	exit(3);
}

#define MATCH(x, str) (strncmp(x, str, sizeof(str) - 1) == 0)
int main(int argc, char** argv) {
	int i, j;
	int fill_flag = 0;
	int fill_byte = 0xff;
	const char* output_file = NULL;
	const char* popt;

	int skip1 = 0;
	int skip2 = 0;

	FILE* output = NULL;

	if (argc < 2)
		usage(argv[0]);

	for (i = 1; i < argc; i++) {
		if (argv[i][0] != '-')
			continue; /* no option but file name */
		if MATCH(argv[i], "--output=") {
			output_file = argv[i] + (sizeof("output=") - 1);
			continue;
		}
		if MATCH(argv[i], "--fillbyte=") {
			fill_flag = 1;
			fill_byte = strtoul(argv[i] + (sizeof("fillbyte=") - 1), NULL, 0);
			continue;
		}
		if MATCH(argv[i], "--fill") {
			fill_flag = 1;
			continue;
		}
		/* no recognized string option, but could be a single letter option */
		/* here, we check for single letter options: o, f, F */
		popt = argv[i];
		for (j = 1; popt[j]; j++) {
			if (popt[j] == 'f') {
				fill_flag = 1;
				continue; /* inner loop, over j */
			}
			if (popt[j] == 'F') {
				if (i + 1 >= argc)
					usage(argv[0]);
				fill_flag = 1;
				fill_byte = strtoul(argv[i + 1], NULL, 0);
				if (skip1 == 0)
					skip1 = i + 1;
				else if (skip2 == 0)
					skip2 = i + 1;
				else
					usage(argv[0]);

				i++;
				continue; /* j loop */
			}
			if (popt[j] == 'o') {
				if (i + 1 >= argc)
					usage(argv[0]);

				output_file = argv[i + 1];
				if (skip1 == 0)
					skip1 = i + 1;
				else if (skip2 == 0)
					skip2 = i + 1;
				else
					usage(argv[0]);

				i++;
				continue; /* j loop */
			}
			/* here: no recognized single letter option */
			usage(argv[0]);
		}
	}

	if (! output_file) {
		while (argv[argc - 1][0] == '-')
			argc--;
		output_file = argv[argc - 1];
		argc--;
	}

	output = fopen(output_file, "w");

	if (! output) {
		printf("Error: cannot open output file %s\n", output_file);
		exit(4);
	}

	/* now do the hard work */
	for (i = 1; i < argc; i++) {
		if (argv[i][0] == '-')
			continue; /* option, no file */
		if (i == skip1 || i == skip2)
			continue;
		load_file(argv[i]);
	}

	if (real_len == 0)
		printf("Warning: 0 bytes were read\n");

	if (fill_flag && fill_byte != 0xff) {
		for (i = 0; i < real_len; i++) {
			if ((data_buf[i] & 0xff00) != 0)
				data_buf[i] = fill_byte;
		}
	}
	
	write_file(output, ! fill_flag);
	fclose(output);
	if (data_buf)
		free(data_buf);

	return 0;
}