#include <merror.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <elf.h>
#include <string.h>
#include <stddef.h>
#include <melf.h>

#include "asm.h"
#include "gen.h"
#include "tab.h"

extern char *current_file;

#define SYMSEC_STUB -1
#define SYMSEC_EXTERN -1

#define SEC_ALIGN 0x1000

static int elf_rel_type(enum reference_type ty) {
	switch (ty) {
	case REF_NONE:
		return R_MIPS_NONE;
	case REF_MIPS_16:
		return R_MIPS_16;
	case REF_MIPS_26:
		return R_MIPS_26;
	case REF_MIPS_PC16:
		return R_MIPS_PC16;
	case REF_MIPS_LO16:
		return R_MIPS_LO16;
	case REF_MIPS_HI16:
		return R_MIPS_HI16;
	}

	return R_MIPS_NONE;
}

static int elf_section_init_reltab(struct section *sec,
				   struct elf_section *elf_sec)
{
	Elf32_Rel *reltab = malloc(sizeof(Elf32_Rel) *
		     sec->reftab.len);

	if (reltab == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}
	for (uint32_t i = 0; i < sec->reftab.len; i++) {
		Elf32_Rel *rel = &reltab[i];
		struct reference *ref = &sec->reftab.references[i];
		rel->r_offset = B32(ref->offset);
		int sym = ref->symbol->tabidx + 1;
		int type = elf_rel_type(ref->type);
		rel->r_info = B32(ELF32_R_INFO(sym, type));
	}

	elf_sec->reltab_len = sec->reftab.len;
	elf_sec->reltab = reltab;

	return M_SUCCESS;
}

static int elf_section_init(struct section *sec, struct elf_section *elf_sec)
{
	elf_sec->data = sec;
	elf_sec->shdr_idx = 0; // dont know yet
	elf_sec->reltab_shidx = 0; // dont know yet
	elf_sec->reltab_len = sec->reftab.len;
	elf_sec->reltab = NULL;

	if (sec->reftab.len && elf_section_init_reltab(sec, elf_sec))
		return M_ERROR;

	return M_SUCCESS;
}

/* free an elf section */
static void elf_section_free(struct elf_section *sec)
{
	if (sec->reltab != NULL)
		free(sec->reltab);
}

static int asm_init_sections(struct assembler *assembler)
{
	struct section *sections = assembler->gen.sections;
	uint32_t len = assembler->gen.sections_len;

	struct elf_section *elftab = malloc(sizeof(struct elf_section) * len);
	if (elftab == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (uint32_t i = 0; i < len; i++) {
		struct elf_section *elfsec = &elftab[i];
		elfsec->data = &sections[i];
		if (elf_section_init(&sections[i], elfsec)) {
			free(elftab);
			return M_ERROR;
		}
	}

	assembler->sections = elftab;
	assembler->section_len = len;
	return M_SUCCESS;
}

static int elf_sym_bind(enum symbol_type ty) {
	switch (ty) {
	case SYM_LOCAL:
		return STB_LOCAL;
	case SYM_GLOBAL:
		return STB_GLOBAL;
	case SYM_EXTERN:
		return STB_GLOBAL;
	}

	return STB_GLOBAL;
}

static int asm_init_symtab(struct assembler *assembler) {
	struct symbol_table *symtab = &assembler->gen.symtab;
	size_t len = symtab->len + 1;
	Elf32_Sym *elftab = malloc(sizeof(Elf32_Sym) * len);
	if (elftab == NULL) {
		PERROR("cannot alloc");
	}

	// add null entry
	elftab[0] = (Elf32_Sym) {0};

	// add rest of the entries
	for (uint32_t i = 0; i < symtab->len; i++) {
		struct symbol *sym = &symtab->symbols[i];
		int bind = elf_sym_bind(sym->type);
		int type = STT_NOTYPE;

		// get name
		size_t str_off;
		if (strtab_write_str(&assembler->strtab, sym->name.str,
		     &str_off)) {
			free(elftab);
			return M_ERROR;
		}

		elftab[i+1] = (Elf32_Sym) {
			.st_name = B32(str_off),
			.st_info = ELF32_ST_INFO(bind, type),
			.st_size = 0,
			.st_other = 0,
			.st_value = B32(sym->offset),
			.st_shndx = 0,
		};
	}

	assembler->symbols = elftab;
	assembler->symtab_len = len;

	return M_SUCCESS;
}

static int parse_file(struct assembler *assembler)
{
	if (generate_mips32r6(&assembler->gen))
		return M_ERROR;
	if (asm_init_sections(assembler))
		return M_ERROR;
	if (asm_init_symtab(assembler))
		return M_ERROR;
	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->section_len; // sections
	max_entries += assembler->section_len; // reltabs per section

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

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

	size_t str_off;
	uint32_t count = 0;

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

	// reltables
	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		const char *prefix = ".reltab";
		char reltab_name[MAX_EXT_LENGTH + strlen(prefix)];

		if (sec->reltab_len == 0)
			continue;

		strcpy(reltab_name, prefix);
		strncat(reltab_name, sec->data->name.str,
				MAX_EXT_LENGTH - strlen(prefix));

		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 = B32(str_off),
			.sh_type = B32(SHT_REL),
			.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_Rel)),
		};
	}

	// for each section
	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		const char *name = sec->data->name.str;

		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 = B32(count);

		shdr[count++] = (Elf32_Shdr){
		    .sh_name = B32(str_off),
		    .sh_type = B32(sec->data->execute ?
					SHT_PROGBITS : SHT_NOBITS),
		    .sh_flags = B32(
				(sec->data->write << 0) |
				(sec->data->execute << 2) |
				SHF_ALLOC),
		    .sh_addr = 0,
		    .sh_offset = 0,
		    .sh_size = 0,
		    .sh_link = 0,
		    .sh_info = 0,
		    .sh_addralign = B32(SEC_ALIGN),
		    .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 = B32(str_off),
		.sh_type = B32(SHT_SYMTAB),
		.sh_flags = 0,
		.sh_addr = 0,
		.sh_offset = 0,
		.sh_size = 0,
		.sh_link = 1,
		.sh_info = 0,
		.sh_addralign = B32(1),
		.sh_entsize = B32(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 = 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,
	};

	// 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 = 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,
	};

	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		if (sec->reltab_len == 0)
			continue;
		shdr[sec->reltab_shidx].sh_link =
			B32(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;
	uint32_t ptr = 0;

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

	// reltbls
	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		if (sec->reltab_len == 0)
			continue;
		int idx = sec->reltab_shidx;
		int len = sec->reltab_len;
		shdr[idx].sh_offset = B32(ptr);
		shdr[idx].sh_size = B32(len * sizeof(Elf32_Rel));
		ptr += len * sizeof(Elf32_Rel);
	}

	// sections
	size_t v_addr = 0;
	for (uint32_t i = 0; i < assembler->section_len; i++) {

		size_t pad = v_addr % SEC_ALIGN;
		if (pad)
			pad = SEC_ALIGN - pad;
		v_addr += pad;

		struct elf_section *sec = &assembler->sections[i];
		uint32_t idx = sec->shdr_idx;
		uint32_t size = sec->data->len;
		shdr[idx].sh_offset = B32(ptr);
		shdr[idx].sh_size = B32(size);
		shdr[idx].sh_addr = B32(v_addr);
		v_addr += size;
		ptr += size;
	}

	// symtab
	{
		uint32_t len = assembler->symtab_len;
		uint32_t size = len * sizeof(Elf32_Sym);
		shdr[assembler->symtab_shidx].sh_offset = B32(ptr);
		shdr[assembler->symtab_shidx].sh_link =
			B32(assembler->strtab_shidx);
		shdr[assembler->symtab_shidx].sh_size = B32(size);
		ptr += size;
	}

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

	// shstrtab
	shdr[assembler->shstrtab_shidx].sh_offset = B32(ptr);
	shdr[assembler->shstrtab_shidx].sh_size =
		B32(assembler->shstrtab.size);
	ptr += assembler->shstrtab.size;
	// shdr
	ehdr->e_shoff = B32(ptr);
}

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

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

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

	// reltbls
	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		void *ptr = sec->reltab;
		int len = sec->reltab_len;
		if (len < 1)
			continue;
		fwrite(ptr, sizeof(Elf32_Rel), len, out);
	}

	// sections
	for (uint32_t i = 0; i < assembler->section_len; i++) {
		struct elf_section *sec = &assembler->sections[i];
		void *ptr = sec->data->data;
		size_t size = sec->data->len;
		fwrite(ptr, 1, size, out);
	}

	// sym tbl
	fwrite(assembler->symbols, sizeof(Elf32_Sym), assembler->symtab_len,
		out);

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

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

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

	// close
	fclose(out);

	return M_SUCCESS;
}

static void update_sym_shndx(struct assembler *assembler)
{
	for (uint32_t i = 1; i < assembler->symtab_len; i++) {
		Elf32_Sym *esym = &assembler->symbols[i];
		struct symbol *sym = &assembler->gen.symtab.symbols[i - 1];

		// get shindx
		int shindx = 0;
		if (sym->secidx != SYM_SEC_STUB)
			shindx = assembler->sections[sym->secidx].shdr_idx;
		else if (sym->type == SYM_EXTERN)
			shindx = 0;

		esym->st_shndx = B16(shindx);
	}
}

static int assemble_elf(struct assembler *assembler, const char *out)
{
	if (assemble_shdr(assembler, &assembler->shdr, &assembler->shdr_len))
		return M_ERROR;

	Elf32_Ehdr ehdr = MIPS_ELF_EHDR;
	ehdr.e_shnum = B16(assembler->shdr_len);
	ehdr.e_shstrndx = B16(assembler->shstrtab_shidx);
	update_offsets(assembler, &ehdr);
	update_sym_shndx(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)
{
	assembler->shdr = NULL;
	assembler->symbols = NULL;
	assembler->sections = NULL;
	assembler->strtab.ptr = NULL;
	assembler->shstrtab.ptr = NULL;
	assembler->gen.sections = NULL;
	assembler->gen.symtab.symbols = NULL;
	assembler->section_len = 0;

	if (generator_init(path, &assembler->gen))
		return M_ERROR;

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

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

	return M_SUCCESS;
}

void assembler_free(struct assembler *assembler)
{
	if (assembler->shdr)
		free(assembler->shdr);
	if (assembler->symbols)
		free(assembler->symbols);
	if (assembler->sections) {
		for (uint32_t i = 0; i < assembler->section_len; i++)
			elf_section_free(&assembler->sections[i]);
		free(assembler->sections);
	}

	strtab_free(&assembler->strtab);
	strtab_free(&assembler->shstrtab);
	generator_free(&assembler->gen);
}