#include <stdlib.h>
#include <merror.h>
#include <melf.h>
#include <mips.h>
#include <mips1.h>
#include <mips32r2.h>
#include <mips32r6.h>

#include "tab.h"
#include "gen.h"
#include "parse.h"

///
/// generation functions
///

static int gen_directive_whb(struct generator *gen, const void *data,
			     uint32_t count, uint32_t len)
{
	// TODO: endianess
	for (uint32_t i = 0; i < count; i++) {
		void *ptr = (char *) data + (len * i);
		if (section_push(gen->current, ptr, len))
			return M_ERROR;
	}

	return M_SUCCESS;
}

static int gen_directive(struct generator *gen,
			 const struct expr *const e)
{
	const struct expr_directive *const expr = &e->directive;
	int res = M_SUCCESS;

	switch (expr->type) {
	case EXPR_DIRECTIVE_ALIGN:
		if (expr->align < 1) {
			ERROR("alignment cannot be zero");
			print_curr_line(gen, e);
			return M_ERROR;
		}
		gen->current->align = expr->align;
		break;
	case EXPR_DIRECTIVE_SPACE:
		res = section_zero(gen->current, expr->space);
		break;
	case EXPR_DIRECTIVE_WORD:
		res = gen_directive_whb(gen, expr->words, expr->len,
			   sizeof(uint32_t));
		break;
	case EXPR_DIRECTIVE_HALF:
		res = gen_directive_whb(gen, expr->halfs, expr->len,
			   sizeof(uint16_t));
		break;
	case EXPR_DIRECTIVE_BYTE:
		res = gen_directive_whb(gen, expr->bytes, expr->len,
			   sizeof(uint8_t));
		break;
	case EXPR_DIRECTIVE_SECTION:
		res = section_get(gen, &gen->current, &expr->section);
		break;
	case EXPR_DIRECTIVE_EXTERN: {
		struct symbol *sym;
		res = symtab_find_or_stub(&gen->symtab, &sym, &expr->label);
		if (res == M_SUCCESS)
			sym->type = SYM_EXTERN;
		break;
	}
	case EXPR_DIRECTIVE_GLOBL: {
		struct symbol *sym;
		res = symtab_find_or_stub(&gen->symtab, &sym, &expr->label);
		if (res == M_SUCCESS)
			sym->type = SYM_GLOBAL;
		break;
	}
	case EXPR_DIRECTIVE_ASCII:
		res = section_push(gen->current, expr->string.str,
		     expr->string.len - 1);
		break;
	case EXPR_DIRECTIVE_ASCIIZ:
		res = section_push(gen->current, expr->string.str,
		     expr->string.len);
		break;
	}

	return res;
}

static int gen_constant(struct generator *gen, struct expr_const *const expr)
{
	(void) gen;
	(void) expr;

	ERROR("constants not yet implemented");
	return M_ERROR;
}

static int gen_ins_write_state(
	struct generator *gen,
	union mips32_instruction ins, // the instruction to modify
	struct gen_ins_state *state,  // the current read state
	char *grammer)		      // the gramemr to parse
{
	// grammer pointer
	char *ptr = grammer;

	// new reference type if needed
	enum reference_type reftype = REF_NONE;

	// hardcoded
	struct gen_ins_override hstate;
	bool hc = false;

	// read values into state
	while (*ptr != '\0') {
		// parse next dsl entry
		size_t skip;
		enum grammer_type gmr = get_gmr_type(ptr, &skip);

		// check for dsl hardcoded register argument
		hc = false;
		if (*(ptr + skip) == '=') {
			// get value pointer
			char *value = ptr + skip + 1;
			hc = true;
			// parse value
			if (get_gmr_hc(state, &hstate, value)) {
				BUG("invalid grammer: %s", grammer);
				exit(1);
			}
		}

		// skip till next comma
		for (;*ptr != '\0' && *ptr != ','; ptr++);
		if (*ptr == ',')
			ptr++;

		switch (gmr) {
		case GMR_RD:
			ins.rd = hc ? hstate.reg : state->rd;
			break;
		case GMR_RS:
			ins.rs = hc ? hstate.reg : state->rs;
			break;
		case GMR_RT:
			ins.rt = hc ? hstate.reg : state->rt;
			break;
                case GMR_FS:
			ins.fs = hc ? hstate.fpreg : state->fs;
			break;
                case GMR_FT:
			ins.ft = hc ? hstate.fpreg : state->ft;
			break;
                case GMR_FD:
			ins.fd = hc ? hstate.fpreg : state->fd;
			break;
		case GMR_IMMD:
			ins.immd = hc ? hstate.immd : state->immd;
			break;
                case GMR_CC:
			ins.cc = state->cc;
			break;
                case GMR_CODE:
			ins.code = state->code;
			break;
		case GMR_POS:
			ins.shamt = state->pos;
			break;
		case GMR_SIZE:
			ins.rd = state->size ? state->size - 1 : 0;
			break;
		case GMR_HB:
			ins.hb = hc ? hstate.immd : state->hb;
			break;
		case GMR_HINT:
			ins.rt = state->hint;
			break;
		case GMR_OFFSET:
			ins.offset = state->offset;
			reftype = REF_MIPS_PC16;
			break;
		case GMR_OFFSET_BASE:
			ins.offset = state->offset;
			ins.rs = state->base;
			reftype = REF_MIPS_16;
			break;
		case GMR_INDEX_BASE:
			ins.rt = state->index;
			ins.rs = state->base;
			break;
		case GMR_TARGET:
			ins.target = state->target;
			reftype = REF_MIPS_26;
			break;
		case GMR_HI:
			ins.immd = state->target >> 16;
			reftype = REF_MIPS_HI16;
			break;
		case GMR_LO:
			ins.immd = state->target & 0x0000FFFF;
			reftype = REF_MIPS_LO16;
			break;
                  break;
                }
        }

	// get offset for reference (if needed)
	uint32_t offset = gen->current->len;
	size_t zeros = offset % gen->current->align;
	if (zeros)
		zeros = gen->current->align - zeros;
	offset += zeros;

	// write instructon to section
	uint32_t raw = B32(ins.raw);
	if (section_push(gen->current, &raw, sizeof(uint32_t))) {
		return M_ERROR;
	}

	// create reference (if needed)
	if (reftype != REF_NONE && state->label != NULL) {
		struct symbol *sym;

		if (symtab_find_or_stub(&gen->symtab, &sym, state->label))
			return M_ERROR;

		struct reference ref = {
			.type = reftype,
			.symbol = sym,
			.offset = offset
		};

		if (reftab_push(&gen->current->reftab, &ref)) {
			return M_ERROR;
		}
	}

	return M_SUCCESS;
}

static int gen_ins(struct generator *gen, struct expr *const expr)
{
	struct mips32_grammer *grammer = NULL;
	for (uint32_t i = 0; i < gen->grammers_len; i++) {
		struct mips32_grammer *temp = &gen->grammers[i];
		if (strcasecmp(temp->name, expr->instruction.name.str) != 0)
			continue;
		grammer = temp;
		break;
	}

	if (grammer == NULL) {
		ERROR("unknown instruction");
		print_curr_line(gen, expr);
		return M_ERROR;
	}

	struct gen_ins_state state;
	state.label = NULL;

	// read in the values from the parser
	if (gen_ins_read_state(gen, expr, &state, grammer))
		return M_ERROR;

	// write the values into the instructions
	// ...and then the sections
	if (grammer->pseudo_len > 0) {
		// write pseudo
		for (int i = 0; i < grammer->pseudo_len; i++) {
			union mips32_instruction ins = gen->instructions[
				grammer->pseudo_grammer[i].enum_index];
			if (gen_ins_write_state(gen, ins, &state,
				grammer->pseudo_grammer[i].update))
				return M_ERROR;
		}
	} else {
		// write real
		union mips32_instruction ins
			= gen->instructions[grammer->enum_index];
		if (gen_ins_write_state(gen, ins, &state, grammer->grammer))
			return M_ERROR;
	}

	return M_SUCCESS;
}

static int gen_label(struct generator *gen, struct string *const label)
{
	uint32_t offset = gen->current->len;
	ptrdiff_t secidx = gen->current - gen->sections;
	size_t zeros = offset % gen->current->align;
	if (zeros)
		zeros = gen->current->align - zeros;
	offset += zeros;

	struct symbol *sym;
	/* update existing symbol (if exists) */
	if (symtab_find(&gen->symtab, &sym, label->str) == M_SUCCESS) {
		if (sym->secidx != SYM_SEC_STUB) {
			// symbols that are not labeled stub are fully defined,
			// it is a error to redefine them
			ERROR("redefined symbol '%s'", label->str);
			return M_ERROR;
		}
		sym->secidx = secidx;
		sym->offset = offset;
	/* create a new symbol */
	} else {
		struct symbol new = {
			.secidx = secidx,
			.offset = offset,
			.type = SYM_LOCAL,
		};
		if (string_clone(&new.name, label))
			return M_ERROR;
		if (symtab_push(&gen->symtab, &new)) {
			string_free(&new.name);
			return M_ERROR;
		}
	}

	return M_SUCCESS;
}

/* run codegen for next expression */
static int generate_next(struct generator *gen)
{
	struct expr expr;
	int res = M_SUCCESS;

	// get the next expression
	if ((res = parser_next(&gen->parser, &expr)))
		return  res;

	// if its not a segment directive
	// (and we dont have a section)
	// create the default
	if ((
		expr.type != EXPR_DIRECTIVE ||
		expr.directive.type != EXPR_DIRECTIVE_SECTION) &&
		gen->current == NULL) {
		// create .data section
		struct string temp = {
			.str = ".data",
			.len = 5,
			.size = 5,
			.allocated = false
		};
		if (section_get(gen, &gen->current, &temp)) {
			expr_free(&expr);
			return M_ERROR;
		}
	}

	res = M_SUCCESS;
	switch (expr.type) {
	case EXPR_DIRECTIVE:
		res = gen_directive(gen, &expr);
		break;
	case EXPR_CONSTANT:
		res = gen_constant(gen, &expr.constant);
		break;
	case EXPR_INS:
		res = gen_ins(gen, &expr);
		break;
	case EXPR_LABEL:
		res = gen_label(gen, &expr.label);
		break;
	}

	expr_free(&expr);
	return res;
}

static int generate(struct generator *gen)
{
	int res;
	while (res = generate_next(gen), 1) {
		if (res == M_ERROR)
			return M_ERROR;
		if (res == M_EOF)
			break;
	}

	return M_SUCCESS;
}

//
// the following functions set default generator
// state values for different versions of the mips
// isa
//

/* run codegen with the mips32r6 specification */
int generate_mips32r6(struct generator *gen)
{
	gen->instructions_len = __MIPS32R6_INS_LEN;
	gen->instructions = mips32r6_instructions;
	gen->grammers_len = __MIPS32R6_GRAMMER_LEN;
	gen->grammers = mips32r6_grammers;
	return generate(gen);
}

/* run codegen with the mips32r2 specification */
int generate_mips32r2(struct generator *gen)
{
	gen->instructions_len = __MIPS32R2_INS_LEN;
	gen->instructions = mips32r2_instructions;
	gen->grammers_len = __MIPS32R2_GRAMMER_LEN;
	gen->grammers = mips32r2_grammers;
	return generate(gen);
}

/* run codegen with the mips32r6 specification */
int generate_mips1(struct generator *gen)
{
	gen->instructions_len = __MIPS1_INS_LEN;
	gen->instructions = mips1_instructions;
	gen->grammers_len = __MIPS1_GRAMMER_LEN;
	gen->grammers = mips1_grammers;
	return generate(gen);
}

//
// constructors and deconstructors
//

int generator_init(const char *file, struct generator *gen)
{
	if (parser_init(file, &gen->parser))
		return M_ERROR;
	if (symtab_init(&gen->symtab))
		return M_ERROR;
	gen->sections = NULL;
	gen->sections_len = 0;
	gen->sections_size = 0;
	return M_SUCCESS;
}

void generator_free(struct generator *gen)
{
	parser_free(&gen->parser);
	symtab_free(&gen->symtab);
	for (size_t i = 0; i < gen->sections_len; i++)
		section_free(&gen->sections[i]);
	free(gen->sections);
}