#include <merror.h>
#include <mips.h>
#include <mips32.h>
#include <stdio.h>
#include <stdlib.h>
#include <elf.h>
#include <string.h>
#include <stddef.h>

#include "asm.h"
#include "mlimits.h"
#include "parse.h"
#include "parse_mips32.h"

extern char *current_file;

static int handle_directive(struct assembler *asm,
			    struct mips32_directive *directive)
{
	switch (directive->type) {
	case MIPS32_DIRECTIVE_SECTION: {
		struct section_table *sec_tbl = &asm->parser.sec_tbl;
		struct section *sec;
		if (sectbl_get(sec_tbl, &sec, directive->name)
				== M_SUCCESS) {
			sec_tbl->current = sec;
			break;
		}

		if (sectbl_alloc(sec_tbl, &sec, directive->name))
			return M_ERROR;

		sec_tbl->current = sec;
		break;
	}

	case MIPS32_DIRECTIVE_ALIGN: {
		asm->parser.sec_tbl.current->alignment =
			1 << directive->align;
		break;
	}

	case MIPS32_DIRECTIVE_SPACE: {
		struct section_entry entry;
		entry.type = ENT_NO_DATA;
		entry.size = directive->space;
		if (sec_push(asm->parser.sec_tbl.current, entry))
				return M_ERROR;
		break;
	}

	case MIPS32_DIRECTIVE_WORD: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_WORD;
			entry.word = directive->words[i];
			if (sec_push(asm->parser.sec_tbl.current, entry))
				return M_ERROR;
		}
		break;
	}

	case MIPS32_DIRECTIVE_HALF: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_HALF;
			entry.half = directive->halfs[i];
			if (sec_push(asm->parser.sec_tbl.current, entry))
				return M_ERROR;
		}
		break;
	}

	case MIPS32_DIRECTIVE_BYTE: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_BYTE;
			entry.byte = directive->bytes[i];
			if (sec_push(asm->parser.sec_tbl.current, entry))
				return M_ERROR;
		}
		break;
	}

        case MIPS32_DIRECTIVE_EXTERN: {
		struct symbol symbol;
		if (symtbl_find(&asm->parser.sym_tbl, NULL, directive->name)
				== M_SUCCESS) {
			ERROR("cannot extern local symbol '%s'",
				directive->name);
			return M_ERROR;
		}

		symbol = (struct symbol) {
			.name = "",
			.sec = asm->parser.sec_tbl.current,
			.index = asm->parser.sec_tbl.current->count,
			.flag = SYM_EXTERNAL,
		};
		strcpy(symbol.name, directive->name);

		if (symtbl_push(&asm->parser.sym_tbl, symbol))
			return M_ERROR;

		break;
	}

        case MIPS32_DIRECTIVE_GLOBL: {
		struct symbol symbol;
		if (symtbl_find(&asm->parser.sym_tbl, NULL, directive->name)
				== M_SUCCESS) {
			symbol.flag = SYM_GLOBAL;
			break;
		}

		symbol = (struct symbol) {
			.name = "",
			.sec = NULL,
			.index = 0,
			.flag = SYM_GLOBAL,
		};
		strcpy(symbol.name, directive->name);

		if (symtbl_push(&asm->parser.sym_tbl, symbol))
			return M_ERROR;

		break;
	}
        }

        return M_SUCCESS;
}

static int parse_file(struct assembler *asm)
{
	struct parser *parser = &asm->parser;

	while (1) {
		struct expr expr;
		if (parser_next(parser, &expr)) {
			break;
		}

		switch (expr.type) {
		case EXPR_INS:
			struct section_entry entry;
			entry.type = ENT_INS;
			entry.size = sizeof(struct mips32_instruction);
			entry.ins = expr.ins;
			if (sec_push(parser->sec_tbl.current, entry))
				return M_ERROR;
			break;

		case EXPR_DIRECTIVE:
			if (handle_directive(asm, &expr.directive.mips32))
				return M_ERROR;
			break;

                case EXPR_CONSTANT:
                case EXPR_LABEL:
			// nothing needed to be done
			break;
                }
        }

	struct section_meta *meta = malloc(sizeof(struct section_meta) *
				    parser->sec_tbl.count);
	if (meta == NULL) {
		ERROR("cannot alloc");
		return M_ERROR;
	}

	asm->meta = meta;

	size_t ptr = 0;
	for (uint32_t i = 0; i < parser->sec_tbl.count; i++) {
		struct section *sec = &parser->sec_tbl.sections[i];
		meta[i].v_addr = ptr;
		ptr += sec_size(sec);
	}

	return M_SUCCESS;
}

static int assemble_phdr(struct assembler *asm, Elf32_Phdr **res,
			 uint32_t *res2)
{
	struct parser *parser = &asm->parser;
	Elf32_Phdr *phdr = malloc(sizeof(Elf32_Phdr) *
				  parser->sec_tbl.count);
	if (phdr == NULL) {
		ERROR("cannot alloc");
		return M_ERROR;;
	}

	for (uint32_t i = 0; i < parser->sec_tbl.count; i++) {
		Elf32_Phdr *hdr = &phdr[i];
		struct section *sec = &parser->sec_tbl.sections[i];
		size_t size = sec_size(sec);
		hdr->p_type = PT_LOAD;
                hdr->p_flags = (sec->execute << 0) |
			(sec->write << 1) |
			(sec->read << 2);
                hdr->p_offset = 0;
		hdr->p_vaddr = 0;
		hdr->p_paddr = 0;
		hdr->p_filesz = size;
		hdr->p_memsz = size;
		hdr->p_align = sec->alignment;
	}

	*res = phdr;
	*res2 = parser->sec_tbl.count;
	return M_SUCCESS;
}

static int assemble_symtbl(struct assembler *asm, Elf32_Sym **res,
			   uint32_t *res2)
{
	Elf32_Sym *stbl = malloc(sizeof(Elf32_Sym) * asm->parser.sym_tbl
				 .count);
	size_t size = 0;

	if (stbl == NULL)
		return M_ERROR;

	for (uint32_t i = 0; i < asm->parser.sym_tbl.count; i++) {
		struct symbol *sym = &asm->parser.sym_tbl.symbols[i];
		size_t str_off;
		unsigned char bind;
		unsigned char type = STT_NOTYPE;

		if (strtbl_write_str(&asm->str_tbl, sym->name, &str_off)) {
			free(stbl);
			return M_ERROR;
		}

		if (sym->flag != SYM_LOCAL)
			bind = STB_LOCAL;
		else
			bind = STB_GLOBAL;

		stbl[i] = (Elf32_Sym) {
			.st_name = str_off,
			.st_value = sym->index,
			.st_size = 0,
			.st_info = ELF32_ST_INFO(bind, type),
			.st_other = ELF32_ST_VISIBILITY(STV_DEFAULT),
			.st_shndx = asm->meta[sym->sec->index].shdr_idx,
		};
		size = i + 1;
	};

	*res = stbl;
	*res2 = size;

	return M_SUCCESS;
}

static int assemble_reltbl_sec(struct assembler *asm, Elf32_Sym *symtbl,
			       uint32_t symtbl_len, struct section *sec)
{
	uint32_t len = 0;

	for (uint32_t i = 0; i < asm->parser.ref_tbl.count; i++) {
		struct reference *ref = &asm->parser.ref_tbl.references[i];
		if (ref->section->index == sec->index) {
			len++;
		}
	}

	if (len == 0) {
		asm->meta[sec->index].reltbl = NULL;
		return M_SUCCESS;
	}

	Elf32_Rela *reltbl = malloc(sizeof(Elf32_Rela) * len);

	if (reltbl == NULL) {
		ERROR("cannot alloc");
		return M_ERROR;
	}

	for (uint32_t i = 0; i < asm->parser.ref_tbl.count; i++) {
		struct reference *ref = &asm->parser.ref_tbl.references[i];
		struct mips32_instruction *ins = &ref->section->
					entries[ref->index].ins.mips32;
		struct section_meta *meta = &asm->meta[ref->section->index];

		if (ref->section->index != sec->index) {
			continue;
		}

		int32_t addend = 0;
		unsigned char type = 0;
		switch (ref->type) {
		case REF_OFFESET:
			addend = ins->B_data.offset;
			type = R_MIPS_PC16;
			break;
		case REF_TARGET:
			addend = ins->J_data.target;
			type = R_MIPS_26;
			break;
		}

		int32_t symidx = -1;

		for (uint32_t i = 0; i < symtbl_len; i++) {
			Elf32_Sym *sym = &symtbl[i];
			const char *str = &asm->str_tbl.ptr[sym->st_name];
			if (strcmp(ref->name, str) == 0) {
				symidx = i;
				break;
			}
		}

		if (symidx == -1) {
			ERROR("undefined symbol '%s'", ref->name);
			free(reltbl);
			return M_ERROR;
		}

		reltbl[i] = (Elf32_Rela) {
			.r_info = ELF32_R_INFO(symidx, type),
			.r_addend = addend,
			.r_offset = meta->v_addr +
				sec_index(ref->section, ref->index),
		};
	};

	asm->meta[sec->index].reltbl_len = len;
	asm->meta[sec->index].reltbl = reltbl;

	return M_SUCCESS;
}

static int assemble_reltbl(struct assembler *asm, Elf32_Sym *symtbl,
			   uint32_t symtbl_len)
{
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		struct section *sec = &asm->parser.sec_tbl.sections[i];
		if (assemble_reltbl_sec(asm, symtbl, symtbl_len, sec))
			return M_ERROR;
	}

	return M_SUCCESS;
}

static int assemble_shdr(struct assembler *asm, Elf32_Shdr **res,
			 uint32_t *res2)
{
	uint32_t max_entries = 4; // symtab, strtab, shstrtab
	max_entries += asm->parser.sec_tbl.count; // sections
	max_entries += asm->parser.sec_tbl.count; // reltabs per section

	Elf32_Shdr *shdr = malloc(sizeof(Elf32_Shdr) * max_entries);

	size_t str_off;
	uint32_t count = 0;

	// eeltables
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {

		if (asm->meta[i].reltbl == NULL)
			continue;

		struct section *sec = &asm->parser.sec_tbl.sections[i];
		const char *prefix = ".reltab.";
		char reltab_name[MAX_LEX_LENGTH + 8];

		strcpy(reltab_name, prefix);
		strcat(reltab_name, sec->name);

		if (strtbl_write_str(&asm->shstr_tbl, reltab_name, &str_off)) {
			free(shdr);
			return M_ERROR;
		}

		asm->meta[i].reltbl_idx = count;
		shdr[count++] = (Elf32_Shdr) {
			.sh_name = str_off,
			.sh_type = SHT_RELA,
			.sh_flags = 0,
			.sh_addr = 0,
			.sh_offset = 0,
			.sh_size = 0,
			.sh_link = 0,
			.sh_info = 0,
			.sh_addralign = 1,
			.sh_entsize = sizeof(Elf32_Rela),
		};
	}

	// for each section
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		struct section *sec = &asm->parser.sec_tbl.sections[i];
		char name[MAX_LEX_LENGTH+1] = ".";
		strcat(name, sec->name);
		if (strtbl_write_str(&asm->shstr_tbl, name, &str_off)) {
			free(shdr);
			return M_ERROR;
		}
		asm->meta[i].shdr_idx = count;
		if (asm->meta[i].reltbl != NULL)
			shdr[asm->meta[i].reltbl_idx].sh_info = count;
                shdr[count++] = (Elf32_Shdr){
                    .sh_name = str_off,
                    .sh_type = SHT_PROGBITS,
                    .sh_flags = (sec->write << 0) | (sec->execute << 2) |
			SHF_ALLOC,
                    .sh_addr = 0,
                    .sh_offset = 0,
                    .sh_size = 0,
                    .sh_link = 0,
                    .sh_info = 0,
                    .sh_addralign = sec->alignment,
                    .sh_entsize = sizeof(struct mips32_instruction),
                };
        }

	// symbol table
	if (strtbl_write_str(&asm->shstr_tbl, ".symtab", &str_off)) {
		free(shdr);
		return M_ERROR;
	}

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

	// string table
	if (strtbl_write_str(&asm->shstr_tbl, ".strtab", &str_off)) {
		free(shdr);
		return M_ERROR;
	}

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

	// sh string table
	if (strtbl_write_str(&asm->shstr_tbl, ".shstrtab", &str_off)) {
		free(shdr);
		return M_ERROR;
	}

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

	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		if (asm->meta[i].reltbl == NULL)
			continue;
		shdr[asm->meta[i].reltbl_idx].sh_link = asm->symtbl_idx;
	}

	*res = shdr;
	*res2 = count;

	return M_SUCCESS;
}

static int assemble_file(struct assembler *asm)
{
	Elf32_Phdr *phdr;
	Elf32_Shdr *shdr;
	Elf32_Sym *symtbl;
	uint32_t phdr_len;
	uint32_t shdr_len;
	uint32_t symtbl_len;

	if (assemble_symtbl(asm, &symtbl, &symtbl_len))
		return M_ERROR;

	if (assemble_reltbl(asm, symtbl, symtbl_len)) {
		free(symtbl);
		return M_ERROR;
	};

	if (assemble_phdr(asm, &phdr, &phdr_len)) {
		free(symtbl);
		return M_ERROR;
	}

	if (assemble_shdr(asm, &shdr, &shdr_len)) {
		free(symtbl);
		free(phdr);
		return M_ERROR;
	};

	Elf32_Ehdr ehdr = {
		.e_ident = {
			[EI_MAG0] = ELFMAG0,
			[EI_MAG1] = ELFMAG1,
			[EI_MAG2] = ELFMAG2,
			[EI_MAG3] = ELFMAG3,
			[EI_CLASS] = ELFCLASS32,
			[EI_DATA] = ELFDATA2LSB,
			[EI_VERSION] = EV_CURRENT,
			[EI_OSABI] = ELFOSABI_NONE,
			[EI_ABIVERSION] = 0x00,
			[EI_PAD] = 0x00,
		},
		.e_type = ET_REL,
		.e_machine = EM_MIPS,
		.e_version = EV_CURRENT,
		.e_entry = 0x00,
		.e_phoff = 0x00,
		.e_shoff = 0x00,
		.e_flags = EF_MIPS_ARCH_32R6,
		.e_ehsize = sizeof(Elf32_Ehdr),
		.e_phentsize = sizeof(Elf32_Phdr),
		.e_phnum = phdr_len,
		.e_shentsize = sizeof(Elf32_Shdr),
		.e_shnum = shdr_len,
		.e_shstrndx = asm->shstrtbl_idx,
	};

	uint32_t ptr = 0;

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

	// phdr
	ehdr.e_phoff = ptr;
	ptr += phdr_len * sizeof(Elf32_Phdr);

	// reltbls
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		if (asm->meta[i].reltbl == NULL)
			continue;
		int idx = asm->meta[i].reltbl_idx;
		int len = asm->meta[i].reltbl_len;
		shdr[idx].sh_offset = ptr;
		shdr[idx].sh_size = len * sizeof(Elf32_Rela);
		ptr += len * sizeof(Elf32_Rela);
	}

	// sections
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		int idx = asm->meta[i].shdr_idx;
		phdr[i].p_offset = ptr;
		phdr[i].p_vaddr = ptr;
		phdr[i].p_paddr = ptr;
		shdr[idx].sh_offset = ptr;
		shdr[idx].sh_size = phdr[i].p_filesz;
		shdr[idx].sh_addr = phdr[i].p_vaddr;
		shdr[idx].sh_addralign = phdr[i].p_align;
		ptr += phdr[i].p_filesz;
	}

	// symtbl
	shdr[asm->symtbl_idx].sh_offset = ptr;
	shdr[asm->symtbl_idx].sh_link = asm->strtbl_idx;
	shdr[asm->symtbl_idx].sh_size = symtbl_len * sizeof(Elf32_Sym);
	ptr += symtbl_len * sizeof(Elf32_Sym);

	// strtbl
	shdr[asm->strtbl_idx].sh_offset = ptr;
	shdr[asm->strtbl_idx].sh_size = asm->str_tbl.size;
	ptr += asm->str_tbl.size;

	// shstrtbl
	shdr[asm->shstrtbl_idx].sh_offset = ptr;
	shdr[asm->shstrtbl_idx].sh_size = asm->shstr_tbl.size;
	ptr += asm->shstr_tbl.size;

	// shdr
	ehdr.e_shoff = ptr;

	FILE *out = fopen("out.o", "w");

	// ehdr
	fwrite(&ehdr, sizeof(Elf32_Ehdr), 1, out);

	// phdr
	fwrite(phdr, sizeof(Elf32_Phdr), phdr_len, out);

	// reltbls
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		if (asm->meta[i].reltbl == NULL)
			continue;
		void *ptr = asm->meta[i].reltbl;
		int len = asm->meta[i].reltbl_len;
		asm->meta[i].reltbl = NULL;
		fwrite(ptr, sizeof(Elf32_Rela), len, out);
	}

	// sections
	for (uint32_t i = 0; i < asm->parser.sec_tbl.count; i++) {
		struct section *sec = &asm->parser.sec_tbl.sections[i];
		for (uint32_t j = 0; j < sec->count; j++) {
			struct section_entry *entry = &sec->entries[j];
			size_t size = entry->size;
			fwrite(&entry->data, size, 1, out);
			while(size % sec->alignment) {
				uint8_t zero = 0;
				fwrite(&zero, 1, 1, out);
				size++;
			}
		}
	}

	// sym tbl
	fwrite(symtbl, sizeof(Elf32_Sym), symtbl_len, out);

	// str tbl
	fwrite(asm->str_tbl.ptr, asm->str_tbl.size, 1, out);

	// shstr tbl
	fwrite(asm->shstr_tbl.ptr, asm->shstr_tbl.size, 1, out);

	// shdr
	fwrite(shdr, sizeof(Elf32_Shdr), shdr_len, out);

	// cleanip
	fclose(out);
	free(shdr);
	free(phdr);
	free(symtbl);

	return M_SUCCESS;
}

int assemble_file_mips32(char *path)
{
	struct assembler asm;
	int res = M_SUCCESS;

	current_file = path;

	if (assembler_init(&asm, path))
		return M_ERROR;

	mips32_parser_init(&asm.parser);

	if (res == M_SUCCESS)
		res = parse_file(&asm);

	if (res == M_SUCCESS)
		res = assemble_file(&asm);

	assembler_free(&asm);

	return res;
}