#include <merror.h>
#include <mips.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <elf.h>
#include <string.h>
#include <stddef.h>
#include <arpa/inet.h>

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

extern char *current_file;

#define SYMSEC_STUB -1
#define SYMSEC_EXTERN -1

#define SEC_ALIGN 0x1000

static int create_symbol(struct assembler *assembler,
			 const char name[MAX_LEX_LENGTH],
			 ssize_t section_idx,
			 size_t section_offset,
			 unsigned char bind)
{
	size_t str_off;
	if (strtab_write_str(&assembler->strtab, name, &str_off))
		return M_ERROR;

	Elf32_Sym symbol = {
		.st_name = htonl(str_off),
		.st_value = htonl(section_offset),
		.st_size = 0,
		.st_info = ELF32_ST_INFO(bind, STT_NOTYPE),
		.st_other = ELF32_ST_VISIBILITY(STV_DEFAULT),
		.st_shndx = htons(section_idx),
	};

	// dont put magic flag values inside symbol, only real indexes
	if (section_idx < 0)
		symbol.st_shndx = 0;

	if (symtab_push(&assembler->symtab, symbol, section_idx))
		return M_ERROR;

	return M_SUCCESS;
}

static int find_symbol_or_stub(struct assembler *assembler,
			       const char name[MAX_LEX_LENGTH],
			       Elf32_Sym **res,
			       size_t *res2)
{
	if (symtab_find(&assembler->symtab, res, res2, name) == M_SUCCESS)
		return M_SUCCESS;

	if (create_symbol(assembler, name, SYMSEC_STUB, 0, STB_LOCAL))
		return M_ERROR;

	size_t idx = assembler->symtab.len - 1;

	if (res != NULL)
		*res = &assembler->symtab.symbols[idx];
	if (res2 != NULL)
		*res2 = idx;

	return M_SUCCESS;
}

static int handle_directive(struct assembler *assembler,
			    struct mips_directive *directive)
{
	switch (directive->type) {
	case MIPS_DIRECTIVE_SECTION: {
		struct section_table *sec_tbl = &assembler->sectab;
		struct section *sec;
		if (sectab_get(sec_tbl, &sec, directive->name)
				== M_SUCCESS) {
			sec_tbl->current = sec;
			break;
		}

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

		sec_tbl->current = sec;
		break;
	}

	case MIPS_DIRECTIVE_ALIGN: {
		assembler->sectab.current->alignment =
			1 << directive->align;
		break;
	}

	case MIPS_DIRECTIVE_SPACE: {
		struct section_entry entry;
		entry.type = ENT_NO_DATA;
		entry.size = directive->space;
		if (sec_push(assembler->sectab.current, entry))
				return M_ERROR;
		break;
	}

	case MIPS_DIRECTIVE_WORD: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_WORD;
			entry.word = directive->words[i];
			entry.size = sizeof(uint32_t);
			if (sec_push(assembler->sectab.current,
				entry))
				return M_ERROR;
		}
		break;
	}

	case MIPS_DIRECTIVE_HALF: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_HALF;
			entry.half = directive->halfs[i];
			entry.size = sizeof(uint16_t);
			if (sec_push(assembler->sectab.current,
				entry))
				return M_ERROR;
		}
		break;
	}

	case MIPS_DIRECTIVE_BYTE: {
		for (uint32_t i = 0; i < directive->len; i++) {
			struct section_entry entry;
			entry.type = ENT_BYTE;
			entry.byte = directive->bytes[i];
			entry.size = sizeof(uint8_t);
			if (sec_push(assembler->sectab.current,
				entry))
				return M_ERROR;
		}
		break;
	}

        case MIPS_DIRECTIVE_EXTERN: {
		if (symtab_find(&assembler->symtab, NULL, NULL,
			directive->name) == M_SUCCESS) {
			ERROR("cannot extern local symbol '%s'",
				directive->name);
			return M_ERROR;
		}

		if (create_symbol(assembler, directive->name, SYMSEC_EXTERN, 0,
		    STB_GLOBAL))
			return M_ERROR;

		break;
	}

        case MIPS_DIRECTIVE_GLOBL: {
		Elf32_Sym *sym;
		if (symtab_find(&assembler->symtab, &sym, NULL,
			directive->name) == M_SUCCESS) {
			sym->st_info = ELF32_ST_INFO(STB_GLOBAL, STT_NOTYPE);
			break;
		}

		if (create_symbol(assembler, directive->name, SYMSEC_STUB, 0,
			STB_GLOBAL))
			return M_ERROR;

		break;
	}

        case MIPS_DIRECTIVE_ASCII: {
		struct section_entry entry;
		entry.type = ENT_STR;
		entry.size = strlen(directive->name);
		memcpy(entry.str, directive->name, entry.size);
		if (sec_push(assembler->sectab.current, entry))
			return M_ERROR;
		break;
	}

        case MIPS_DIRECTIVE_ASCIIZ: {
		struct section_entry entry;
		entry.type = ENT_STR;
		entry.size = strlen(directive->name) + 1;
		memcpy(entry.str, directive->name, entry.size);
		if (sec_push(assembler->sectab.current, entry))
			return M_ERROR;
		break;
	}
        }

        return M_SUCCESS;
}

static int handle_label(struct assembler *assembler,
			const char name[MAX_LEX_LENGTH])
{
	struct section *cur = assembler->sectab.current;

	Elf32_Sym *ref;
	size_t symidx;

	if (symtab_find(&assembler->symtab, &ref, &symidx, name) == M_SUCCESS) {
		ssize_t *sec = &assembler->symtab.sections[symidx];

		// check if the symbol is acutally jus a stub, if so
		// we need to update it
		if (*sec == SYMSEC_STUB) {
			*sec = cur->index;
			ref->st_value = htonl(sec_size(cur));
			return M_SUCCESS;
		}

		ERROR("redefined symbol '%s'", name);
		return M_ERROR;
	}

	if (create_symbol(assembler, name, cur->index, sec_size(cur),
		   STB_LOCAL))
		return M_ERROR;

	return M_SUCCESS;
}

static int handle_ins(struct assembler *assembler,
		      struct ins_expr *expr)
{
	struct section *sec = assembler->sectab.current;
	size_t secidx = sec->len;

	for (size_t i = 0; i < expr->ins_len; i++) {
		union mips_instruction_data *ins =
				&expr->ins[i].data;
		struct reference *ref =
				&expr->ref[i];
		struct section_entry entry;

		entry.type = ENT_INS;
		entry.size = sizeof(union mips_instruction_data);
		entry.ins = htonl(ins->raw);

		if (sec_push(sec, entry))
			return M_ERROR;

		unsigned char type = 0;
		switch (ref->type) {
		case REF_NONE:
			continue;
		case REF_OFFESET:
			type = R_MIPS_PC16;
			break;
		case REF_TARGET:
			type = R_MIPS_26;
			break;
		}

		size_t symidx;
		if (find_symbol_or_stub(assembler, ref->name, NULL, &symidx))
			return M_ERROR;

		Elf32_Rela rel = {
			.r_info = htonl(ELF32_R_INFO(symidx, type)),
			.r_addend = htonl(ref->addend),
			.r_offset = htonl(sec_index(sec, secidx + i)),
		};

		if (reltab_push(&sec->reltab, rel))
			return M_ERROR;

		break;
	}

	return M_SUCCESS;
}

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

	while (1) {
		struct expr expr;
		int res = parser_next(parser, &expr);

		if (res == M_ERROR)
			return M_ERROR;

		if (res == M_EOF)
			return M_SUCCESS;

		switch (expr.type) {
		case EXPR_INS:
			if (handle_ins(assembler, &expr.ins))
				return M_ERROR;
			break;
		case EXPR_DIRECTIVE:
			if (handle_directive(assembler,
				&expr.directive))
				return M_ERROR;
			break;

                case EXPR_LABEL:
			if (handle_label(assembler, expr.label))
				return M_ERROR;
			break;

                case EXPR_CONSTANT:
			break;
		}
	}

	return M_SUCCESS;
}

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

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

	*res = phdr;
	*res2 = assembler->sectab.len;
	return M_SUCCESS;
}

static int assemble_shdr(struct assembler *assembler, Elf32_Shdr **res,
			 uint32_t *res2)
{
	uint32_t max_entries = 0;
	max_entries += 1; // null
	max_entries += 1; // symtab
	max_entries += 1; // strtab
	max_entries += 1; // shtrtab
	max_entries += assembler->sectab.len; // sections
	max_entries += assembler->sectab.len; // reltabs per section

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

	size_t str_off;
	uint32_t count = 0;

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

	// reltables
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		const char *prefix = ".reltab.";
		char reltab_name[MAX_LEX_LENGTH + 8];

		if (sec->reltab.len == 0)
			continue;

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

		if (strtab_write_str(&assembler->shstrtab,
		       reltab_name, &str_off)) {
			free(shdr);
			return M_ERROR;
		}

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

	// for each section
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		char name[MAX_LEX_LENGTH+1] = ".";

		strcat(name, sec->name);
		if (strtab_write_str(&assembler->shstrtab, name, &str_off)) {
			free(shdr);
			return M_ERROR;
		}

		sec->shdr_idx = count;
		if (sec->reltab.len != 0)
			shdr[sec->reltab_shidx].sh_info = htonl(count);

                shdr[count++] = (Elf32_Shdr){
                    .sh_name = htonl(str_off),
                    .sh_type = htonl(SHT_PROGBITS),
                    .sh_flags = htonl(
				(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 = htonl(sec->alignment),
                    .sh_entsize = 0,
                };
        }

	// symbol table
	if (strtab_write_str(&assembler->shstrtab, ".symtab", &str_off)) {
		free(shdr);
		return M_ERROR;
	}

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

	// string table
	if (strtab_write_str(&assembler->shstrtab, ".strtab", &str_off)) {
		free(shdr);
		return M_ERROR;
	}

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

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

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

	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		if (sec->reltab.len == 0)
			continue;
		shdr[sec->reltab_shidx].sh_link =
			htonl(assembler->symtab_shidx);
	}

	*res = shdr;
	*res2 = count;

	return M_SUCCESS;
}

static void update_offsets(struct assembler *assembler, Elf32_Ehdr *ehdr)
{
	Elf32_Shdr *shdr = (Elf32_Shdr *) assembler->shdr;
	Elf32_Phdr *phdr = (Elf32_Phdr *) assembler->phdr;
	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 = htonl(ptr);
	ptr += assembler->phdr_len * sizeof(Elf32_Phdr);

	// reltbls
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		if (sec->reltab.len == 0)
			continue;
		int idx = sec->reltab_shidx;
		int len = sec->reltab.len;
		shdr[idx].sh_offset = htonl(ptr);
		shdr[idx].sh_size = htonl(len * sizeof(Elf32_Rela));
		ptr += len * sizeof(Elf32_Rela);
	}

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

	// sections
	size_t v_addr = 0x00400000;
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		uint32_t idx = sec->shdr_idx;
		uint32_t size = ntohl(phdr[i].p_filesz);
		phdr[i].p_offset = htonl(ptr);
		phdr[i].p_vaddr = htonl(v_addr);
		phdr[i].p_paddr = htonl(v_addr);
		shdr[idx].sh_offset = htonl(ptr);
		shdr[idx].sh_size = phdr[i].p_filesz;
		shdr[idx].sh_addr = phdr[i].p_vaddr;
		v_addr += size;
		ptr += size;
	}

	// symtab
	shdr[assembler->symtab_shidx].sh_offset = htonl(ptr);
	shdr[assembler->symtab_shidx].sh_link = htonl(assembler->strtab_shidx);
	shdr[assembler->symtab_shidx].sh_size =
		htonl(assembler->symtab.len * sizeof(Elf32_Sym));
	ptr += assembler->symtab.len * sizeof(Elf32_Sym);

	// strtab
	shdr[assembler->strtab_shidx].sh_offset = htonl(ptr);
	shdr[assembler->strtab_shidx].sh_size = htonl(assembler->strtab.size);
	ptr += assembler->strtab.size;

	// shstrtab
	shdr[assembler->shstrtab_shidx].sh_offset = htonl(ptr);
	shdr[assembler->shstrtab_shidx].sh_size =
		htonl(assembler->shstrtab.size);
	ptr += assembler->shstrtab.size;

	// shdr
	ehdr->e_shoff = htonl(ptr);
}

static void update_sym_shindx(struct assembler *assembler)
{
	for (size_t i = 0; i < assembler->symtab.len; i++) {
		Elf32_Sym *sym = &assembler->symtab.symbols[i];
		ssize_t sec = assembler->symtab.sections[i];

		if (sec >= 0) {
			sym->st_shndx = htons(assembler->
				sectab.sections[sec].shdr_idx);
		}
	}
}

static int write_file(struct assembler *assembler, Elf32_Ehdr *ehdr,
		       const char *path)
{
	FILE *out = fopen(path, "w");

	if (out == NULL) {
		ERROR("cannot write '%s'", path);
		return M_ERROR;
	}

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

	// phdr
	fwrite(assembler->phdr, sizeof(Elf32_Phdr), assembler->phdr_len, out);

	// reltbls
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		if (sec->reltab.len == 0)
			continue;
		void *ptr = sec->reltab.data;
		int len = sec->reltab.len;
		fwrite(ptr, sizeof(Elf32_Rela), len, out);
	}

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

	// sections
	for (uint32_t i = 0; i < assembler->sectab.len; i++) {
		struct section *sec = &assembler->sectab.sections[i];
		for (uint32_t j = 0; j < sec->len; 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(assembler->symtab.symbols, sizeof(Elf32_Sym),
		assembler->symtab.len, out);

	// str tbl
	fwrite(assembler->strtab.ptr, assembler->strtab.size, 1, out);

	// shstr tbl
	fwrite(assembler->shstrtab.ptr, assembler->shstrtab.size, 1, out);

	// shdr
	fwrite(assembler->shdr, sizeof(Elf32_Shdr), assembler->shdr_len, out);

	fclose(out);

	return M_SUCCESS;
}

static int assemble_elf(struct assembler *assembler, const char *out)
{
	if (assemble_phdr(assembler, (Elf32_Phdr **) &assembler->phdr,
		   &assembler->phdr_len)) {
		return M_ERROR;
	}

	if (assemble_shdr(assembler, (Elf32_Shdr **) &assembler->shdr,
		   &assembler->shdr_len)) {
		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] = ELFDATA2MSB,
			[EI_VERSION] = EV_CURRENT,
			[EI_OSABI] = ELFOSABI_NONE,
			[EI_ABIVERSION] = 0x00,
			[EI_PAD] = 0x00,
		},
		.e_type = htons(ET_REL),
		.e_machine = htons(EM_MIPS),
		.e_version = htonl(EV_CURRENT),
		.e_entry = 0x00,
		.e_phoff = 0x00,
		.e_shoff = 0x00,
		.e_flags = htonl(EF_MIPS_ARCH_32R6),
		.e_ehsize = htons(sizeof(Elf32_Ehdr)),
		.e_phentsize = htons(sizeof(Elf32_Phdr)),
		.e_phnum = htons(assembler->phdr_len),
		.e_shentsize = htons(sizeof(Elf32_Shdr)),
		.e_shnum = htons(assembler->shdr_len),
		.e_shstrndx = htons(assembler->shstrtab_shidx),
	};

	update_offsets(assembler, &ehdr);
	update_sym_shindx(assembler);

	if (write_file(assembler, &ehdr, out))
		return M_ERROR;

	return M_SUCCESS;
}

int assemble_file(struct assembler_arguments args)
{
	struct assembler assembler;
	int res = M_SUCCESS;

	current_file = args.in_file;

	if (assembler_init(&assembler, args.in_file))
		return M_ERROR;

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

	if (res == M_SUCCESS)
		res = assemble_elf(&assembler, args.out_file);

	assembler_free(&assembler);

	return res;
}

int assembler_init(struct assembler *assembler, const char *path)
{
	if (lexer_init(path, &assembler->lexer))
		return M_ERROR;

	if (parser_init(&assembler->lexer, &assembler->parser))
		return M_ERROR;

	if (strtab_init(&assembler->shstrtab))
		return M_ERROR;

	if (strtab_init(&assembler->strtab))
		return M_ERROR;

	if (symtab_init(&assembler->symtab))
		return M_ERROR;

	if (sectab_init(&assembler->sectab))
		return M_ERROR;

	assembler->symtab.strtab = &assembler->strtab;
	assembler->phdr = NULL;
	assembler->shdr = NULL;

	return M_SUCCESS;
}

void assembler_free(struct assembler *assembler)
{
	if (assembler->phdr)
		free(assembler->phdr);
	if (assembler->shdr)
		free(assembler->shdr);

	sectab_free(&assembler->sectab);
	symtab_free(&assembler->symtab);
	strtab_free(&assembler->strtab);
	strtab_free(&assembler->shstrtab);

	parser_free(&assembler->parser);
	lexer_free(&assembler->lexer);
}