#include "irm.h"

#include <assert.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>

#include <backend.h>
#include <x86.h>

#define STACK_SUB(irm, syment, type) \
	do { \
	        irm->stackdepth += sizeof(type); \			
                syment->loc = irm->stackdepth; \
                syment->len = sizeof(type); \
    while(0)

// TODO: Add scopes.

void irm_init(struct irm *irm)
{
	assert(irm != NULL);
	irm->storage = malloc(1024 * 1024); // TODO: platform-agnostic
	irm->size = (1024 * 1024) / 8; 
	irm->curr_pos = 0;
	irm->state = IRM_GRACEFUL;
	symtab_init(&irm->symtab);
}

void irm_panic(struct irm *irm)
{
	irm->state = IRM_PANIC;
}

void irm_stmt_enter_scope(struct irm *irm)
{
    assert(irm != NULL);
    // Append 1 level.
}

void irm_stmt_leave_scope(struct irm *irm)
{
    assert(irm != NULL);	
}

void irm_stmt_var_decl(struct irm *irm, size_t *offset)
{
 	uint64_t* mem = irm->storage;

	uint64_t ltype = mem[*offset];
	(*offset)++;

	uint64_t symbol_type = mem[*offset];
	if(symbol_type != SYM_VAR) {
		DIE("Symbol type invalid! sym:'%x'", symbol_type);
	}
	(*offset)++;
	uint64_t symid = mem[*offset];
	
	if(irm->symtab.count < symid) {
		DIE("Symbol id: '%d' not found! This is a compiler bug.\n",
		    symid);
	}

	struct symbol_entry *syment = &irm->symtab.entries[symid];
	irm->stackdepth += sizeof(uint64_t);
	syment->loc = irm->stackdepth;
	
    DEBUG("Symbol id: %d\n", symid);
	(*offset)++;	
	(*offset)++;
	switch(ltype) {
	case VAR_DECL_QWORD: {
		uint64_t value = mem[*offset];
		syment->len = sizeof(uint64_t);
		(*offset)++;
		x86_mov_r_i64(RAX, value);
		x86_push_r64(RAX);
		break;
	}			
	case VAR_DECL_DWORD: {
		uint64_t value = mem[*offset];
		syment->len = sizeof(uint32_t);
		(*offset)++;
		x86_push_i32(value);
		break;
	}
	case VAR_DECL_WORD: {
		uint64_t value = mem[*offset];
		syment->len = sizeof(uint32_t);
		(*offset)++;
		x86_push_i32(value); // TODO: Replace with immediate 16 instruction
		break;
	}
	case VAR_DECL_BYTE: {		
		uint64_t value = mem[*offset];
		(*offset)++;
		syment->len = sizeof(uint16_t); 
		x86_push_i8(value);
		break;
	}
	default:
		DIE("Invalid type supplied! This is a compiler bug. type: %x.\n"
		    , ltype);
		break;
	}
}

void irm_stmt_var_assign(struct irm *irm, size_t *offset)
{
	assert(irm != NULL);
	if(offset == NULL) {
		DIE("offset must not be null!");
	}

	(*offset)++;
	uint64_t sym_id = irm->storage[*offset];
	if(sym_id != SYM_VAR) {
		DIE("Variable assignment only supports a symbol of type variable!");
	}

        (*offset)++;
	uint64_t symtab_id = irm->storage[*offset];
	struct symbol_entry *ent = &irm->symtab.entries[symtab_id];
	if(ent->symtype != SYM_VAR) {
		DIE("Symbol '%.*s' is not of type variable!", ent->name.len, ent->name.buf);
	}
	(*offset)++;
	uint64_t value_type = irm->storage[*offset];
	switch(value_type) {
	case NUMBER64: { // FIX STACK ISSUE
		DIE("TODO:: STACK ALGINMENT ISSUE!!");
		(*offset)++;
		uint64_t value = irm->storage[*offset];
	        DEBUG("Assigned nr64 literal %d to '%.*s'.\n", value, ent->name.len, ent->name.buf);
		(*offset)++;
		// TODO: mov [rsp - X], value
		// mov rax, rsp 
		emit8(0x48);
		emit8(0x8B); // MR
		emit8(0xC4); // rsp -> rax
		// sub rax, 8
		emit8(0x48);
		emit8(0x2d);
		emit32(ent->loc);
		// mov rbx, value
		emit8(0x48);
		emit8(0xBB); // b8 + rd
		emit64(value); // FALSE VALUE!! FIX!
		// mov qword ptr [rax], rbx
		emit8(0x48);
		emit8(0x89); // MR
		emit8(0x18); // rbx -> [rax]
		break;
	}
	default:
		DIE("Variable assignment only supports a number literal!");
	}
}

void irm_stmt_done(struct irm *irm)
{
	if(irm->state == IRM_PANIC) return;
        size_t offset = 0;
	
	while(offset < irm->curr_pos - 1) {
		uint64_t type = irm->storage[offset];	
		type &= 0xFFFF0000;

		switch(type) {
		case VAR_DECL:
			irm_stmt_var_decl(irm, &offset);
			break;
		case VAR_ASSIGN:
			irm_stmt_var_assign(irm, &offset);
			break;
		default:
			DIE("Statement identifier not found: %d.", type);
			break;
		}
		DEBUG("Offset: %d, Size: %d\n", offset, irm->curr_pos);
	}
	if(irm->storage[offset] != STMT_DONE) {
		DIE("Statement done was not received! This is a compiler bug.");
	}
	memset(irm->storage, 0, irm->curr_pos);
	irm->curr_pos = 0;
}

void irm_push(struct irm *irm, uint64_t type)
{
	assert(irm != NULL);
	uint64_t *mem = irm->storage + irm->curr_pos;
	mem[0] = type;
	irm->curr_pos += 1;
        DEBUG("added %x to irm.\n", type);
}

void irm_push_64v(struct irm *irm, uint64_t type, uint64_t value)
{
	assert(irm != NULL);
	if(irm->curr_pos + 2 > irm->size)
	        DIE("IRM OVERFLOW!");
	uint64_t* mem = irm->storage + irm->curr_pos;
	mem[0] = type;
	mem[1] = value;
	irm->curr_pos += 2;
        DEBUG("added %d to irm with value %d.\n", type, value);
}

void irm_push_lookup_sym(struct irm *irm, uint64_t type, struct string_view sym)
{
	assert(irm != NULL);

	uint64_t *mem = irm->storage + irm->curr_pos;
	mem[0] = type;
        struct symlookup_result ent = symtab_lookup(&irm->symtab, sym);

	if(ent.ent == NULL) {
		DIE("'%.*s' not found!", sym.len, sym.buf);
	}

	mem[1] = ent.i;
	irm->curr_pos+=2;
}

// TODO:
// 1. Push SYM_TYPE
// 2. Push SYMTAB offset (after adding the sym to it)
void irm_push_sym(struct irm *irm, uint64_t type, struct string_view sym)
{
	assert(irm != NULL);	
	struct symbol_entry entry;
	entry.name = sym;
	entry.loc = 0;
	entry.symtype = type;
	if(symtab_lookup(&irm->symtab, sym).ent != NULL) {
		DIE("Symbol '%.*s' already exists.\n", sym.len, sym.buf);
	}
	size_t symindex = symtab_append(&irm->symtab, entry);
	uint64_t* mem = irm->storage + irm->curr_pos;
	mem[0] = type;
	mem[1] = symindex;
	irm->curr_pos += 2;
	DEBUG("added %.*s to symbol table.\n", sym.len, sym.buf);
}