#include <elf.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <merror.h>
#include <string.h>
#include <sys/stat.h>
#include <melf.h>
#include <fcntl.h>

#include "link.h"
#include "mips.h"

#define SEC_ALIGN 0x1000

static int load_objects(struct linker *linker)
{
	linker->objects = malloc(sizeof(struct object) *
			  linker->args->in_count);
	linker->obj_len = 0;

	if (linker->objects == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (int i = 0; i < linker->args->in_count; i++) {
		char *path = linker->args->in_files[i];
		struct object *obj = &linker->objects[linker->obj_len];

		// check for duplicate
		for (size_t j = 0; j < linker->obj_len; j++) {
			const char *dupname = linker->objects[j].name;
			struct stat st, st2;
			if (stat(path, &st) || stat(dupname, &st2))
				continue;
			if (st.st_ino == st2.st_ino)
				goto skip_obj;
		}

		// load obj file
		linker->obj_len++;
		if (object_load(obj, path))
			return M_ERROR;
skip_obj:
	}

	return M_SUCCESS;
}

/**
 * Relocates all segments with the given name
 * (since they need to be next to eachother)
 */
static int relocate_segment_name(struct linker *linker, const char *name)
{
	for (size_t i = 0; i < linker->obj_len; i++) {
		struct object *obj = &linker->objects[i];
		for (size_t j = 0; j < obj->segment_len; j++) {
			struct segment *seg = &obj->segments[j];

			// check if the segment has already been relocated
			if (seg->new_vaddr != 0)
				continue;

			// make sure the segments name matches what
			// we are looking for
			if (strcmp(seg->name, name) != 0)
				continue;

			if (ADDR_CHK(linker->off, seg->size, UINT32_MAX)) {
				ERROR("linker offset overflow");
				return M_ERROR;
			}

			// is the segment a TEXT type or DATA type??
			if (B32(seg->phdr->p_flags) & PF_X) {
				// TEXT
				if (ADDR_CHK(linker->text_vaddr, seg->size,
							DATA_VADDR_MIN)) {
					ERROR("linker text vaddr overflow");
					return M_ERROR;
				}

				seg->new_off = linker->off;
				seg->new_vaddr = linker->text_vaddr;
				linker->off += seg->size;
				linker->text_vaddr += seg->size;
			} else {
				// DATA
				if (ADDR_CHK(linker->data_vaddr, seg->size,
							UINT32_MAX)) {
					ERROR("linker data vaddr overflow");
					return M_ERROR;
				}

				seg->new_off = linker->off;
				seg->new_vaddr = linker->data_vaddr;
				linker->off += seg->size;
				linker->data_vaddr += seg->size;
			}

			// if this is an existing segment, append this
			// part
			struct segment_table_entry *ent;
			if (segtab_get(&linker->segments, &ent, name) ==
				M_SUCCESS) {
				if (segtab_ent_push(ent, seg))
					return M_ERROR;
			} else {
				// else create a new segment
				if (segtab_push(&linker->segments, NULL, seg))
					return M_ERROR;
			}
		}
	}

	return M_SUCCESS;
}

static int relocate_segments(struct linker *linker)
{
	for (size_t i = 0; i < linker->obj_len; i++) {
		struct object *obj = &linker->objects[i];
		for (size_t j = 0; j < obj->segment_len; j++) {
			struct segment *seg = &obj->segments[j];

			// check if the segment has already been relocated
			if (seg->new_vaddr != 0)
				continue;
			if(relocate_segment_name(linker, seg->name))
				return M_ERROR;
		}
	}

	return M_SUCCESS;
}

static int relocate_symbol(struct linker *linker, struct object *obj,
			   struct symbol_table *symtab, const Elf32_Sym *sym)
{
	size_t shndx = B16(sym->st_shndx);
	if (shndx == 0)
		return M_SUCCESS; // ignore this symbol

	// find the given section
	const char *name = symtab->strtab->data + B32(sym->st_name);

	if (shndx >= obj->shdr_len) {
		ERROR("shdr entry [%d] name out of bounds", shndx);
		return M_ERROR;
	}

	if (B32(sym->st_name) >= symtab->strtab->len) {
		ERROR("symbol name out of bounds");
		return M_ERROR;
	}

	Elf32_Shdr const *shdr = &obj->shdr[shndx];
	struct segment *sec = NULL;
	for (size_t i = 0; i < obj->phdr_len; i++) {
		Elf32_Phdr *temp = &obj->phdr[i];
		if (shdr->sh_offset == temp->p_offset) {
			sec = &obj->segments[i];
			break;
		}
	}

	if (sec == NULL) {
		ERROR("could not locate segment for symbol '%s'", name);
		return M_ERROR;
	}

	struct segment_table_entry *ent = NULL;
	if (segtab_get(&linker->segments, &ent, sec->name)) {
		ERROR("could not locate segment for symbol '%s'", name);
		return M_ERROR;
	}

	// segments start at shindx 1
	ptrdiff_t new_shndx = (ent - linker->segments.entries) + 1;

	size_t str_off = 0;
	if (strtab_push(linker->symtab.strtab, name, &str_off))
		return M_ERROR;

	int32_t off = sec->new_vaddr + B32(sym->st_value);
	Elf32_Sym new = *sym;
	new.st_name = B32(str_off);
	new.st_value = B32(off);
	new.st_shndx = B16(new_shndx);
	new.st_size = 0;

	if (symtab_get(&linker->symtab, NULL, name) == M_SUCCESS) {
		ERROR("cannot link doubly defiend symbol '%s'", name);
		return M_ERROR;
	}

	if (symtab_push(&linker->symtab, &new))
		return M_ERROR;

	return M_SUCCESS;
}

static int relocate_symbols(struct linker *linker)
{
	for (size_t i = 0; i < linker->obj_len; i++) {
		struct object *obj = &linker->objects[i];
		for (size_t j = 0; j < obj->shdr_len; j++) {
			struct symbol_table *symtab = &obj->symtabs[j];
			if (symtab->len < 1)
				continue;

			for (size_t k = 0; k < symtab->len; k++) {
				const Elf32_Sym *sym = &symtab->syms[k];
				if (relocate_symbol(linker, obj, symtab, sym))
					return M_ERROR;
			}
		}
	}
	return M_SUCCESS;
}

static int assemble_phdr(struct linker *linker)
{
	Elf32_Phdr *phdr = malloc(sizeof(Elf32_Phdr) * linker->segments.len);

	if (phdr == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (uint32_t i = 0; i < linker->segments.len; i++) {
		Elf32_Phdr *hdr = &phdr[i];
		struct segment_table_entry *ent = &linker->segments.entries[i];
		size_t size = segtab_ent_size(ent);
		hdr->p_type = B32(PT_LOAD);
		hdr->p_flags = B32(
			(ent->parts[0]->execute << 0) |
			(ent->parts[0]->write << 1) |
			(ent->parts[0]->read << 2));
		hdr->p_offset = B32(ent->off);
		hdr->p_vaddr = B32(ent->vaddr);
		hdr->p_paddr = B32(ent->vaddr);
		hdr->p_filesz = B32(size);
		hdr->p_memsz = B32(size);
		hdr->p_align = B32(SEC_ALIGN);
	}

	linker->phdr = phdr;
	linker->phdr_len = linker->segments.len;
	return M_SUCCESS;
}

static int assemble_shdr(struct linker *linker)
{
	uint32_t max_entries = 0;
	max_entries += 1; // null
	max_entries += 1; // symtab
	max_entries += 1; // strtab
	max_entries += 1; // shstrtab
	max_entries += linker->segments.len; // segments

	Elf32_Shdr *shdr = malloc(sizeof(Elf32_Shdr) * max_entries);
	if (shdr == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}
	linker->shdr = shdr;

	size_t str_off;
	uint32_t count = 0;

	// null
	shdr[count++] = (Elf32_Shdr) {0};

	// segments
	for (uint32_t i = 0; i < linker->segments.len; i++) {
		struct segment_table_entry *ent = &linker->segments.entries[i];
		if (strtab_push(&linker->shstrtab, ent->name, &str_off))
			return M_ERROR;

		shdr[count++] = (Elf32_Shdr) {
			.sh_name = B32(str_off),
			.sh_type = B32(SHT_PROGBITS),
			.sh_flags = B32(
				(ent->parts[0]->write << 0) |
				(ent->parts[0]->execute << 2) |
				SHF_ALLOC),
			.sh_addr = B32(ent->vaddr),
			.sh_offset = B32(ent->off),
			.sh_size = B32(segtab_ent_size(ent)),
			.sh_link = 0,
			.sh_info = 0,
			.sh_addralign = B32(ent->parts[0]->align),
			.sh_entsize = 0,
		};
	}

	// symbol table
	if (strtab_push(&linker->shstrtab, ".symtab", &str_off))
		return M_ERROR;

	linker->symtab_shidx = count;
	shdr[count++] = (Elf32_Shdr) {
		.sh_name = B32(str_off),
		.sh_type = B32(SHT_SYMTAB),
		.sh_flags = 0,
		.sh_addr = 0,
		.sh_offset = 0,
		.sh_size = 0,
		.sh_link = 0,
		.sh_info = 0,
		.sh_addralign = B32(1),
		.sh_entsize = B32(sizeof(Elf32_Sym)),
	};

	// string table
	if (strtab_push(&linker->shstrtab, ".strtab", &str_off))
		return M_ERROR;

	linker->strtab_shidx = count;
	shdr[count++] = (Elf32_Shdr) {
		.sh_name = B32(str_off),
		.sh_type = B32(SHT_STRTAB),
		.sh_flags = B32(SHF_STRINGS),
		.sh_addr = 0,
		.sh_offset = 0,
		.sh_size = 0,
		.sh_link = 0,
		.sh_info = 0,
		.sh_addralign = B32(1),
		.sh_entsize = 0,
	};

	// shstring table
	if (strtab_push(&linker->shstrtab, ".shstrtab", &str_off))
		return M_ERROR;

	linker->shstrtab_shidx = count;
	shdr[count++] = (Elf32_Shdr) {
		.sh_name = B32(str_off),
		.sh_type = B32(SHT_STRTAB),
		.sh_flags = B32(SHF_STRINGS),
		.sh_addr = 0,
		.sh_offset = 0,
		.sh_size = 0,
		.sh_link = 0,
		.sh_info = 0,
		.sh_addralign = B32(1),
		.sh_entsize = 0,
	};

	linker->shdr_len = count;
	return M_SUCCESS;
}

static int relocate_instruction_rela(struct linker *linker,
				     struct segment *seg,
				     Elf32_Rela *rel)
{
	/// get start of the segment part in bytes and the byte offset
	/// of this relocation in that segment part
	uint32_t off = B32(rel->r_offset);
	if (off > seg->size) {
		ERROR("relocation in segment '%s' out of bounds", seg->name);
		return M_ERROR;
	}

	if (B32(rel->r_info) == 0) {
		WARNING("skiping empty relocation entry");
		return M_SUCCESS;
	}

	/// read the relocation entry
	uint8_t idx = B32(rel->r_info) >> 8;
	uint8_t typ = B32(rel->r_info) & 0xFF;
	int32_t add = B32(rel->r_addend);

	/// read the symbol from the relocation
	struct symbol_table *symtab = seg->reltab.symtab;
	if (idx >= symtab->len) {
		ERROR("relocation in segment '%s', symbol index [%d] out of "
			"bounds", seg->name, idx);
		return M_ERROR;
	}
	Elf32_Sym *sym = &symtab->syms[idx];
	const char *symname = symtab->strtab->data + B32(sym->st_name);


	/// get the section header that the symbol is related to
	Elf32_Shdr *shdr = NULL;
	if (B16(sym->st_shndx) >= seg->obj->shdr_len) {
		ERROR("shndx index [%d] out of bounds", B16(sym->st_shndx));
		return M_ERROR;
	}
	shdr = &seg->obj->shdr[B16(sym->st_shndx)];

	/// get the segment that the symbol is in
	struct segment_table_entry *ent;
	const char *segname = seg->obj->shstrtab->data + B32(shdr->sh_name);
	if (B32(shdr->sh_name) >= seg->obj->shstrtab->len) {
		ERROR("relocation segment name out of bounds");
		return M_ERROR;
	}
	if (segtab_get(&linker->segments, &ent, segname)) {
		ERROR("could not locate segment for relocation");
		return M_ERROR;
	}

	uint32_t sym_vaddr = B32(sym->st_value) + ent->vaddr;
	uint32_t *ins_raw = (uint32_t *) &seg->bytes[off];

	union mips_instruction_data ins;
	ins.raw = B32(*ins_raw);

	uint32_t ins_vaddr = seg->new_vaddr + off;

	uint32_t vaddr_abs = sym_vaddr + add;
	 int32_t vaddr_rel = (vaddr_abs - ins_vaddr - 4) >> 2;
	bool warn = false;

	switch (typ) {
	case R_MIPS_16:
		// 16bit absolute
		if (vaddr_abs > UINT16_MAX)
			warn = true;
		ins.immd = (uint16_t)vaddr_abs;
		break;
	case R_MIPS_PC16:
		// 16bit relative shifted
		if (vaddr_rel > INT16_MAX || vaddr_rel < INT16_MIN)
			warn = true;
		ins.offset = vaddr_rel;
		break;
	case R_MIPS_26:
		// 26bit absolute shifted
		if (vaddr_abs >= (1 << 25))
			warn = true;
		ins.target = (vaddr_abs & 0x0FFFFFFF) >> 2;
		break;
	case R_MIPS_PC26_S2:
		// 26bit relative shifted
		if (vaddr_rel >= (1 << 24) || -vaddr_rel > (1 << 24))
			warn = true;
		ins.offs26 = vaddr_rel;
		break;
	case R_MIPS_LO16:
		// lo 16bit absolute
		ins.immd = (uint16_t)(vaddr_abs & 0xFFFF);
		break;
	case R_MIPS_HI16:
		// hi 16bit absolute
		ins.immd = (uint16_t)(vaddr_abs >> 16);
		break;
	default:
		ERROR("do not know how do handle relocation type [%d]", typ);
		return M_ERROR;
	}

	*ins_raw = B32(ins.raw);

	if (warn)
		WARNING("truncating relocation for symbol '%s'", symname);

	return M_SUCCESS;

}

static int relocate_instruction_rel(struct linker *linker,
				    struct segment *seg,
				    Elf32_Rel *rel)
{
	Elf32_Rela temp;
	temp.r_info = rel->r_info;
	temp.r_offset = rel->r_offset;
	temp.r_addend = 0;

	return relocate_instruction_rela(linker, seg, &temp);
}

static int relocate_segment_instructions(struct linker *linker,
					 struct segment *seg)
{
	for (uint32_t i = 0; i < seg->reltab.len; i++) {
		int res = M_SUCCESS;
		if (seg->reltab.type == SHT_RELA)
			res = relocate_instruction_rela(linker, seg,
				   &seg->reltab.rela[i]);
		else if (seg->reltab.type == SHT_REL)
			res = relocate_instruction_rel(linker, seg,
				   &seg->reltab.rel[i]);
		else {
			ERROR("unknown reltab type");
			return M_ERROR;
		}
		if (res)
			return M_ERROR;
	}
	return M_SUCCESS;
}

static int relocate_instructions(struct linker *linker)
{
	for (uint32_t i = 0; i < linker->segments.len; i++) {
		struct segment_table_entry *ent = &linker->segments.entries[i];
		for (uint32_t j = 0; j < ent->len; j++) {
			struct segment *seg = ent->parts[j];
			if (relocate_segment_instructions(linker, seg))
				return M_ERROR;
		}
	}
	return M_SUCCESS;
}

static void update_offsets(struct linker *linker)
{
	uint32_t ptr = 0;

	// we must now correct offsets and sizes in side the ehdr, phdr,
	// and shdr
	ptr += sizeof(Elf32_Ehdr);

	// phdr
	linker->ehdr.e_phoff = B32(ptr);
	ptr += linker->phdr_len * sizeof(Elf32_Phdr);

	// section padding
	{
		uint32_t mod = ptr % SEC_ALIGN;
		if (mod != 0)
			linker->secalign = (SEC_ALIGN - mod);
		else
			 linker->secalign = 0;
		ptr += linker->secalign;
	}

	// sections
	for (uint32_t i = 0; i < linker->segments.len; i++) {
		struct segment_table_entry *ent = &linker->segments.entries[i];
		uint32_t idx = i + 1;
		uint32_t size = segtab_ent_size(ent);
		linker->phdr[i].p_offset = B32(ptr);
		linker->shdr[idx].sh_offset = linker->phdr[i].p_offset;
		ptr += size;
	}

	// symtab
	Elf32_Shdr *symtab = &linker->shdr[linker->symtab_shidx];
	symtab->sh_offset = B32(ptr);
	symtab->sh_link = B32(linker->strtab_shidx);
	symtab->sh_size = B32(linker->symtab.len * sizeof(Elf32_Sym));
	ptr += B32(symtab->sh_size);

	// strtab
	Elf32_Shdr *strtab = &linker->shdr[linker->strtab_shidx];
	strtab->sh_offset = B32(ptr);
	strtab->sh_size = B32(linker->strtab.len);
	ptr += linker->strtab.len;

	// shstrtab
	Elf32_Shdr *shstrtab = &linker->shdr[linker->shstrtab_shidx];
	shstrtab->sh_offset = B32(ptr);
	shstrtab->sh_size = B32(linker->shstrtab.len);
	ptr += linker->shstrtab.len;

	// shdr
	linker->ehdr.e_shoff = B32(ptr);
}

static int write_file(struct linker *linker)
{
	extern char *current_file;
	current_file = linker->args->out_file;

	int fd = open(linker->args->out_file, O_RDWR | O_CREAT, 0711);
	if (fd < 0) {
		PERROR("cannot write");
		return M_ERROR;
	}

	FILE *out = fdopen(fd, "w");
	if (out == NULL) {
		PERROR("cannot write");
		return M_ERROR;
	}

	int res = 0;

	// ehdr
	res |= fwrite(&linker->ehdr, sizeof(Elf32_Ehdr), 1, out);

	// phdr
	res |= fwrite(linker->phdr, sizeof(Elf32_Phdr), linker->phdr_len, out);

	// section padding
	for (uint32_t i = 0; i < linker->secalign; i++) {
		uint8_t zero = 0;
		res |= fwrite(&zero, 1, 1, out);
	}

	// sections
	for (uint32_t i = 0; i < linker->segments.len; i++) {
		struct segment_table_entry *ent = &linker->segments.entries[i];
		for (uint32_t j = 0; j < ent->len; j++) {
			struct segment *seg = ent->parts[j];
			res |= fwrite(seg->bytes, 1, seg->size, out);
		}
	}

	// sym tbl
	res |= fwrite(linker->symtab.syms, sizeof(Elf32_Sym), linker->symtab.len, out);

	// str tbl
	res |= fwrite(linker->strtab.data, 1, linker->strtab.len, out);

	// shstr tbl
	res |= fwrite(linker->shstrtab.data, 1, linker->shstrtab.len, out);

	// shdr
	res |= fwrite(linker->shdr, sizeof(Elf32_Shdr), linker->shdr_len, out);

	if (res < 0) {
		ERROR("cannot write data");
		return M_ERROR;
	}

	fclose(out);

	return M_SUCCESS;
}

static int link_executable(struct linker *linker)
{
	Elf32_Ehdr *ehdr = &linker->ehdr;
	*ehdr = MIPS_ELF_EHDR;

	if (assemble_phdr(linker))
		return M_ERROR;

	if (assemble_shdr(linker))
		return M_ERROR;

	ehdr->e_type = B16(ET_EXEC);
	ehdr->e_phnum = B16(linker->phdr_len);
	ehdr->e_shnum = B16(linker->shdr_len);
	ehdr->e_shstrndx = B16(linker->shstrtab_shidx);

	update_offsets(linker);

	if (write_file(linker))
		return M_ERROR;

	return M_SUCCESS;
}

static int linker_init(struct linker *linker, struct linker_arguments *args)
{
	linker->args = args;
	linker->off = 0;
	linker->text_vaddr = TEXT_VADDR_MIN;
	linker->data_vaddr = DATA_VADDR_MIN;
	linker->objects = NULL;
	linker->segments.size = 0;
	linker->symtab.syms = NULL;
	linker->shstrtab.data = NULL;
	linker->strtab.data = NULL;
	linker->shdr = NULL;
	linker->phdr = NULL;
	if (segtab_init(&linker->segments))
		return M_ERROR;
	if (strtab_init(&linker->shstrtab))
		return M_ERROR;
	if (strtab_init(&linker->strtab))
		return M_ERROR;
	if (symtab_init(&linker->symtab))
		return M_ERROR;
	linker->symtab.strtab = &linker->strtab;
	return M_SUCCESS;
}

static void linker_free(struct linker *linker)
{
	if (linker->objects != NULL) {
		for (size_t i = 0; i < linker->obj_len; i++)
			object_free(&linker->objects[i]);
		free(linker->objects);
	}
	if (linker->shdr != NULL)
		free(linker->shdr);
	if (linker->phdr != NULL)
		free(linker->phdr);
	segtab_free(&linker->segments);
	strtab_free(&linker->shstrtab);
	strtab_free(&linker->strtab);
	symtab_free(&linker->symtab);
}

int link_files(struct linker_arguments args) {
	struct linker linker;

	int res = M_SUCCESS;
	if (res == M_SUCCESS)
		res = linker_init(&linker, &args);
	if (res == M_SUCCESS)
		res = load_objects(&linker);
	if (res == M_SUCCESS)
		res = relocate_segments(&linker);
	if (res == M_SUCCESS)
		res = relocate_symbols(&linker);
	if (res == M_SUCCESS)
		res = relocate_instructions(&linker);
	if (res == M_SUCCESS)
		res = link_executable(&linker);

	linker_free(&linker);
	return res;
}