#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <elf.h>

#include "cyimg.h"

static const uint32_t i2ceep_sizes[] = {
	4096,
	8192,
	16384,
	32768,
	65536, // also Atmel 128k and ST 256k
	131072, // Microchip (except 24LC1026)
};
static const uint32_t i2c_speeds_khz[] = { 100, 400, 1000 };
static const uint32_t spi_speeds_mhz[] = { 10, 20, 30 };

static bool validate_cyimg(size_t size, const uint8_t* data) {
	if (size < 4) return false;

	uint16_t wSignature = data[0] | (((uint16_t)data[1]) << 8);

	if (wSignature != CYIMG_MAGIC) return false;

	if (data[3] == CYIMG_IMAGETYPE_USBIDS)
		return size == 8; // nothing else to verify

	if (data[3] != CYIMG_IMAGETYPE_REGULAR) return false;

	size_t cpos = 4;
	printf("size=0x%zx\n", size);
	while (true) {
		//printf("cpos=0x%zx\n", cpos);
		if (cpos + 8 > size) return false;

		struct cyimg_section sec;
		sec.dLength  = CYIMG_R32(&data[cpos])*4;
		cpos += 4;
		sec.dAddress = CYIMG_R32(&data[cpos]);
		cpos += 4;

		printf("len=0x%x addr=0x%x\n", sec.dLength, sec.dAddress);

		if (sec.dAddress & 3) return false;

		if (sec.dLength == 0) break;

		cpos += sec.dLength;
	}

	//printf("cpos+4=0x%zx size=0x%zx\n", cpos+4, size);
	return true;//cpos + 4 == size; // checksum
}
static bool validate_checksum(size_t size, const uint8_t* data) {
	if (size < 4) return false;

	if (data[3] == CYIMG_IMAGETYPE_USBIDS)
		return true; // nothing else to verify
	if (data[3] != CYIMG_IMAGETYPE_REGULAR) return false;

	size_t cpos = 4;
	uint32_t ckcalc = CYIMG_CHECKSUM_INIT;
	while (true) {
		if (cpos + 8 > size) return false;

		struct cyimg_section sec;
		sec.dLength  = CYIMG_R32(&data[cpos])*4;
		cpos += 4;
		sec.dAddress = CYIMG_R32(&data[cpos]);
		cpos += 4;

		if (sec.dLength == 0) break;

		ckcalc = cyimg_checksum_calc(ckcalc, sec.dLength, &data[cpos]);
		cpos += sec.dLength;
	}

	if (cpos + 4 != size) return false;

	return ckcalc == CYIMG_R32(&data[cpos]);
}

static int write_usbids(FILE* out, size_t size, const uint8_t* data) {
	size_t outsize = sizeof(Elf32_Ehdr) + sizeof(Elf32_Phdr) + 2*(sizeof(Elf32_Nhdr)+4+4);
	uint8_t* buf = calloc(1, outsize);

	Elf32_Ehdr* ehdr = (Elf32_Ehdr*)&buf[0];
	Elf32_Phdr* phdr = (Elf32_Phdr*)&ehdr[1];
	Elf32_Nhdr* nhdr = (Elf32_Nhdr*)&phdr[1];
	uint8_t* ndata = (uint8_t*)&nhdr[1];
	Elf32_Nhdr* nhdr2 = (Elf32_Nhdr*)&ndata[8];
	uint8_t* ndata2 = (uint8_t*)&nhdr2[1];

	ehdr->e_ident[0] = 0x7f;
	ehdr->e_ident[1] = 'E';
	ehdr->e_ident[2] = 'L';
	ehdr->e_ident[3] = 'F';
	ehdr->e_ident[4] = ELFCLASS32;
//#ifdef __LITTLE_ENDIAN__
	ehdr->e_ident[5] = ELFDATA2LSB;
//#else
//	ehdr->e_ident[5] = ELFDATA2MSB;
//#endif
	ehdr->e_ident[6] = EV_CURRENT;
	ehdr->e_ident[7] = ELFOSABI_SYSV;

	ehdr->e_type = ET_EXEC;
	ehdr->e_machine = EM_ARM;
	ehdr->e_version = EV_CURRENT;
	ehdr->e_phoff = sizeof(Elf32_Ehdr);
	ehdr->e_phentsize = sizeof(Elf32_Phdr);
	ehdr->e_phnum = 1;

	phdr->p_type = PT_NOTE;
	phdr->p_offset = sizeof(Elf32_Ehdr)+sizeof(Elf32_Phdr);
	phdr->p_vaddr = sizeof(Elf32_Ehdr)+sizeof(Elf32_Phdr);
	phdr->p_paddr = sizeof(Elf32_Ehdr)+sizeof(Elf32_Phdr);
	phdr->p_filesz = 2*sizeof(Elf32_Nhdr)+2*4+4+2;
	phdr->p_memsz = 2*sizeof(Elf32_Nhdr)+2*4+4+2;
	phdr->p_flags = PF_R;
	phdr->p_align = 4;

	nhdr->n_namesz = 4;
	nhdr->n_descsz = 4;
	nhdr->n_type = ELF_FX3_NOTE_TYPE_USBVIDPID;

	ndata[0] = 'F'; ndata[1] = 'X'; ndata[2] = '3'; ndata[3] = '\0';
	ndata[4] = data[4]; ndata[5] = data[5]; ndata[6] = data[6]; ndata[7] = data[7];

	nhdr2->n_namesz = 4;
	nhdr2->n_descsz = 2;
	nhdr2->n_type = ELF_FX3_NOTE_TYPE_IMAGECTLTYPE;

	ndata2[0] = 'F'; ndata2[1] = 'X'; ndata2[2] = '3'; ndata2[3] = '\0';
	ndata2[4] = data[2]; ndata2[5] = data[3];

	if (fwrite(buf, 1, outsize, out) != outsize) {
		fprintf(stderr, "E: cannot write to output.\n");
		return 1;
	}

	return 0;
}

static int write_regular(FILE* out, size_t size, const uint8_t* data) {
	uint32_t nsections = 0,
			 totalsize = 0,
			 entrypt = 0;

	uint32_t curoff = 4;
	while (true) {
		uint32_t len = CYIMG_R32(&data[curoff])*4;
		curoff += 4;
		uint32_t addr = CYIMG_R32(&data[curoff]);
		curoff += 4;
		if (len == 0) {
			entrypt = addr;
			break;
		}

		printf("section: addr=0x%x off=0x%x size=0x%x\n", addr, curoff, len);
		totalsize += len;
		curoff += len;
		++nsections;
	}

	printf("totalsize=0x%lx nsections=0x%lx\n", totalsize, nsections);
	struct { uint32_t addr, off, size; } insections[nsections];

	curoff = 4;
	for (size_t i = 0; i < nsections; ++i) {
		uint32_t len = CYIMG_R32(&data[curoff])*4;
		curoff += 4;
		uint32_t addr = CYIMG_R32(&data[curoff]);
		curoff += 4;
		if (len == 0) break;

		printf("insection %zu: addr=0x%x off=0x%x size=0x%x\n", i, addr, curoff, len);
		insections[i].addr = addr;
		insections[i].off = curoff;
		insections[i].size = len;

		curoff += len;
	}

	// TODO: allocate data for:
	// * EHDR
	// * n+1 PHDRs
	// * n sections
	// * 1 note for bImageCTL/bImageType
	// * data
	size_t outsize = sizeof(Elf32_Ehdr) + sizeof(Elf32_Phdr)*(1+nsections)
		+ totalsize + (sizeof(Elf32_Nhdr)+4+2 +2);
	uint8_t* buf = calloc(1, outsize);

	Elf32_Ehdr* ehdr = (Elf32_Ehdr*)(buf + 0);
	Elf32_Phdr* phdr = (Elf32_Phdr*)(buf + sizeof(Elf32_Ehdr));
	/*uint8_t* sdata = buf + sizeof(Elf32_Ehdr) + (nsections+1)*sizeof(Elf32_Phdr);
	Elf32_Nhdr* nhdr = (Elf32_Nhdr*)(sdata + totalsize);
	uint8_t* ndata = sdata + totalsize + sizeof(Elf32_Nhdr);*/
	Elf32_Nhdr* nhdr = (Elf32_Nhdr*)(buf + sizeof(Elf32_Ehdr) + (nsections+1)*sizeof(Elf32_Phdr));
	uint8_t* ndata = buf + sizeof(Elf32_Ehdr) + (nsections+1)*sizeof(Elf32_Phdr) + sizeof(Elf32_Nhdr);
	uint8_t* sdata = ndata + 8;

	ehdr->e_ident[0] = 0x7f;
	ehdr->e_ident[1] = 'E';
	ehdr->e_ident[2] = 'L';
	ehdr->e_ident[3] = 'F';
	ehdr->e_ident[4] = ELFCLASS32;
//#ifdef __LITTLE_ENDIAN__
	ehdr->e_ident[5] = ELFDATA2LSB;
//#else
//	ehdr->e_ident[5] = ELFDATA2MSB;
//#endif
	ehdr->e_ident[6] = EV_CURRENT;
	ehdr->e_ident[7] = ELFOSABI_SYSV;

	ehdr->e_type = ET_EXEC;
	ehdr->e_machine = EM_ARM;
	ehdr->e_version = EV_CURRENT;
	ehdr->e_entry = entrypt;
	ehdr->e_phoff = sizeof(Elf32_Ehdr);
	ehdr->e_phentsize = sizeof(Elf32_Phdr);
	ehdr->e_phnum = 1 + nsections;

	for (size_t i = 0, coff = 0; i < nsections; ++i) {
		phdr[i].p_type = PT_LOAD;
		phdr[i].p_offset = coff + (size_t)sdata - (size_t)ehdr;
		phdr[i].p_vaddr = insections[i].addr;
		phdr[i].p_paddr = insections[i].addr;
		phdr[i].p_filesz = insections[i].size;
		phdr[i].p_memsz = insections[i].size;
		phdr[i].p_flags = PF_R|PF_W|PF_X;
		phdr[i].p_align = 4;

		printf("dest=%p src=%p, size=0x%x\n", sdata+coff, data+insections[i].off, insections[i].size);
		memcpy(sdata + coff, data + insections[i].off, insections[i].size);
		coff += insections[i].size;
	}
	phdr[nsections].p_type = PT_NOTE;
	phdr[nsections].p_offset = (size_t)nhdr - (size_t)ehdr;
	phdr[nsections].p_vaddr = (size_t)nhdr - (size_t)ehdr;
	phdr[nsections].p_paddr = (size_t)nhdr - (size_t)ehdr;
	phdr[nsections].p_filesz = sizeof(Elf32_Nhdr)+4+2;
	phdr[nsections].p_memsz = sizeof(Elf32_Nhdr)+4+2;
	phdr[nsections].p_flags = PF_R;
	phdr[nsections].p_align = 4;

	nhdr->n_namesz = 4;
	nhdr->n_descsz = 2;
	nhdr->n_type = ELF_FX3_NOTE_TYPE_IMAGECTLTYPE;

	ndata[0] = 'F'; ndata[1] = 'X'; ndata[2] = '3'; ndata[3] = '\0';
	ndata[4] = data[2]; ndata[5] = data[3];

	if (fwrite(buf, 1, outsize, out) != outsize) {
		fprintf(stderr, "E: cannot write to output.\n");
		return 1;
	}

	return 0;
}

int main(int argc, char* argv[]) {
	if (argc != 2 && argc != 3) {
		fprintf(stderr, "Usage: %s <input file.img> [<output file.elf>]\n", argv[0]);
		return 1;
	}

	FILE* f = fopen(argv[1], "rb");
	if (!f) {
		fprintf(stderr, "E: cannot open file '%s'\n", argv[1]);
		return 1;
	}

	struct stat st;
	if (fstat(fileno(f), &st)) {
		fprintf(stderr, "E: cannot stat file '%s'\n", argv[1]);
		return 1;
	}

	uint8_t* data = (uint8_t*)calloc(1, st.st_size);
	if (fread(data, 1, st.st_size, f) != st.st_size) {
		fprintf(stderr, "E: read error.\n");
		return 1;
	}

	fclose(f);

	if (!validate_cyimg(st.st_size, data)) {
		fprintf(stderr, "E: file does not contain a valid Cypress FX3 image.\n");
		return 1;
	}
	if (!validate_checksum(st.st_size, data)) {
		fprintf(stderr, "W: invalid checksum.\n");
	}

	f = stdout;
	if (argc == 3) {
		f = fopen(argv[2], "wb");
		if (!f) {
			fprintf(stderr, "E: cannot open output file '%s'\n", argv[2]);
			return 1;
		}
	}

	int r = 1;
	if (data[3] == CYIMG_IMAGETYPE_USBIDS) {
		r = write_usbids(f, st.st_size, data);
	} else if (data[3] == CYIMG_IMAGETYPE_REGULAR) {
		r = write_regular(f, st.st_size, data);
	}

	if (argc == 3) fclose(f);
	free(data);
	return 0;
}