#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;

#define SHDR_SYMTBL   0
#define SHDR_STRTBL   1
#define SHDR_SECTIONS 2

static int parse_file(struct parser *parser)
{
	while (1) {
		struct expr expr;
		if (parser_next(parser, &expr)) {
			break;
		}

                if (expr.type == EXPR_INS)
                        if (sectbl_push(&parser->sec_tbl,
                                        parser->sec_tbl.current, expr.ins))
                                return M_ERROR;
	}

        for (uint32_t i = 0; i < parser->ref_tbl.count; i++) {
                struct reference *ref = &parser->ref_tbl.references[i];
                struct symbol *sym;
                struct mips32_instruction *ins;

                if (symtbl_find(&parser->sym_tbl, &sym, ref->name)) {
                        ERROR("undefined symbol '%s'", ref->name);
                        return M_ERROR;
                }

                ins = &ref->section->ins[ref->index].mips32;

                switch (ref->type) {
                case REF_OFFESET:
                        ins->B_data.offset += sym->position -
                                        (ref->section->start + ref->index);
                        break;
                case REF_TARGET:
                        ins->J_data.target += sym->position;
                        break;
                }
        };

        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);
        size_t ins_sz = sizeof(struct mips32_instruction);

        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];

                hdr->p_type = PT_LOAD;
                hdr->p_flags = PF_X | PF_W | PF_R; // FIXME: this is bad
                hdr->p_offset = sec->start * ins_sz;
                hdr->p_vaddr = sec->start * ins_sz;
                hdr->p_paddr = 0x00;
                hdr->p_filesz = sec->count * ins_sz;
                hdr->p_memsz = sec->count * ins_sz;
                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);

        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;

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

                int viz = STB_LOCAL;
                switch (sym->flag) {
                case SYM_LOCAL:
                        viz = STB_LOCAL;
                        break;
                case SYM_GLOBAL:
                case SYM_EXTERNAL:
                        viz = STB_GLOBAL;
                        break;
                }

                stbl[i] = (Elf32_Sym) {
                        .st_name = str_off,
                        .st_value = sym->position,
                        .st_size = 0,
                        .st_info = (unsigned char)
                                        ELF32_ST_INFO(SYMINFO_BT_SELF,
                                                 SYMINFO_FLG_DIRECT),
                        .st_other = (unsigned char)
                                        ELF32_ST_VISIBILITY(viz),
                        .st_shndx = 0, // FIXME: specify section
                };
        };

        *res = stbl;
        *res2 = asm->parser->sym_tbl.count;

        return M_SUCCESS;
}

static int assemble_shdr(struct assembler *asm, Elf32_Shdr **res,
                         uint32_t *res2)
{
        uint32_t entries = 2; // str table and sym tabel
        entries += asm->parser->sec_tbl.count; // sections

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

        size_t str_off;
        if (strtbl_write_str(&asm->str_tbl, ".shsymtab", &str_off)) {
                free(shdr);
                return M_ERROR;
        }

        // symbol table
        shdr[SHDR_SYMTBL] = (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),
        };

        if (strtbl_write_str(&asm->str_tbl, ".shstrtab", &str_off)) {
                free(shdr);
                return M_ERROR;
        }

        // string table
        shdr[SHDR_STRTBL] = (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 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->str_tbl, name, &str_off)) {
                        free(shdr);
                        return M_ERROR;
                }
                shdr[i+SHDR_SECTIONS] = (Elf32_Shdr) {
                        .sh_name = str_off,
                        .sh_type = SHT_PROGBITS,
                        .sh_flags = SHF_WRITE | SHF_ALLOC | SHF_EXECINSTR,
                        .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),
                };
        }

        *res = shdr;
        *res2 = entries;

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

        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);

        // sections
        for (uint32_t i = 0; i < asm->parser->sec_tbl.count; i++) {
                phdr[i].p_offset = ptr;
                phdr[i].p_vaddr = ptr;
                shdr[i+SHDR_SECTIONS].sh_offset = ptr;
                shdr[i+SHDR_SECTIONS].sh_size = phdr[i].p_filesz;
                ptr += phdr[i].p_filesz;
        }

        // symtbl
        shdr[SHDR_SYMTBL].sh_offset = ptr;
        shdr[SHDR_SYMTBL].sh_size = symtbl_len * sizeof(Elf32_Sym);
        ptr += symtbl_len * sizeof(Elf32_Sym);

        // strtbl
        shdr[SHDR_STRTBL].sh_offset = ptr;
        shdr[SHDR_STRTBL].sh_size = asm->str_tbl.size;
        ptr += asm->str_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);

        // 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 mips32_instruction *ins = &sec->ins[j].mips32;
                        fwrite(ins, sizeof(struct mips32_instruction),
                               1, out);
                }
        }

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

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

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

        fclose(out);

        free(shdr);
        free(phdr);
        free(symtbl);

        return M_SUCCESS;
}

int assemble_file_mips32(char *path)
{
	struct lexer lexer;
	struct parser parser;
	current_file = path;
        int res = M_SUCCESS;

	if (lexer_init(current_file, &lexer))
		return M_ERROR;

	if (mips32_parser_init(&lexer, &parser))
		return M_ERROR;

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

        struct assembler assembler;
        assembler.parser = &parser;
        strtbl_init(&assembler.str_tbl);

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

        strtbl_free(&assembler.str_tbl);
	lexer_free(&lexer);
	parser_free(&parser);

        return res;
}