#include "compiler.h"
#include "interpreter.h"
#include "types.h"
#include "include/estr.h"
#include <stdint.h>
#include <inttypes.h>

VariableTable createVariableTable() {
    VariableTable vt;
    vt.capacity = 16;
    vt.count = 0;
    vt.vars = malloc(sizeof(VariableInfo) * vt.capacity);
    return vt;
} 

void addVtVariable(VariableTable* vt, const char* name) {
    // check if it already exists
    for (size_t i = 0; i < vt->count; i++) {
        if (strcmp(vt->vars[i].name, name) == 0) {
            return;
        }
    }

    // increase capacity
    if (vt->count >= vt->capacity) {
        vt->capacity *= 2;
        vt->vars = realloc(vt->vars, sizeof(VariableInfo) * vt->capacity);
    }

    // add the variable
    strcpy(vt->vars[vt->count].name, name);
    vt->vars[vt->count].offset = vt->count * 8;
    vt->count++;
}

int getVariablePos(VariableTable* vt, const char* name) {
    for (size_t i = 0; i < vt->count; i++) {
        if (strcmp(vt->vars[i].name, name) == 0) {
            return vt->vars[i].offset;
        }
    }
    return -1;
}

char* processValueString(GroundArg arg) {
    if (arg.type == VALREF) {
        char* buf = malloc(sizeof(char) * 260);
        snprintf(buf, sizeof(char) * 260, "[%s]", arg.value.refName);
        return buf;
    }
    if (arg.type == VALUE) {
        if (arg.value.value.type != INT) {
            printf("Only int is supported right now\n");
            exit(1);
        }
        char* buf = malloc(sizeof(char) * 64);
        snprintf(buf, sizeof(char) * 260, "%" PRId64, arg.value.value.data.intVal);
        return buf;
    }
    return NULL;
}

char* compileGroundProgram(GroundProgram* program) {
    Estr start = CREATE_ESTR("global _start\nsection .text\n_start:\n");
    Estr data = CREATE_ESTR("section .bss\n");
    Estr helpers = CREATE_ESTR("");

    APPEND_ESTR(helpers, "\n; Helper: Print integer in rax\n");
    APPEND_ESTR(helpers, "print_int:\n");
    APPEND_ESTR(helpers, "  push rbp\n");
    APPEND_ESTR(helpers, "  mov rbp, rsp\n");
    APPEND_ESTR(helpers, "  sub rsp, 32          ; Allocate buffer on stack\n");
    APPEND_ESTR(helpers, "  mov rdi, rsp         ; RDI = buffer pointer\n");
    APPEND_ESTR(helpers, "  add rdi, 31          ; Point to end of buffer\n");
    APPEND_ESTR(helpers, "  mov byte [rdi], 0    ; Null terminator\n");
    APPEND_ESTR(helpers, "  dec rdi\n");
    APPEND_ESTR(helpers, "  mov rbx, 10          ; Divisor\n");
    APPEND_ESTR(helpers, "  test rax, rax\n");
    APPEND_ESTR(helpers, "  jns .positive\n");
    APPEND_ESTR(helpers, "  neg rax              ; Make positive\n");
    APPEND_ESTR(helpers, "  push rax\n");
    APPEND_ESTR(helpers, "  mov byte [rdi], '-'\n");
    APPEND_ESTR(helpers, "  dec rdi\n");
    APPEND_ESTR(helpers, "  pop rax\n");
    APPEND_ESTR(helpers, ".positive:\n");
    APPEND_ESTR(helpers, "  xor rcx, rcx         ; Digit counter\n");
    APPEND_ESTR(helpers, ".convert_loop:\n");
    APPEND_ESTR(helpers, "  xor rdx, rdx\n");
    APPEND_ESTR(helpers, "  div rbx              ; rax = rax/10, rdx = remainder\n");
    APPEND_ESTR(helpers, "  add dl, '0'          ; Convert to ASCII\n");
    APPEND_ESTR(helpers, "  mov [rdi], dl\n");
    APPEND_ESTR(helpers, "  dec rdi\n");
    APPEND_ESTR(helpers, "  inc rcx\n");
    APPEND_ESTR(helpers, "  test rax, rax\n");
    APPEND_ESTR(helpers, "  jnz .convert_loop\n");
    APPEND_ESTR(helpers, "  inc rdi              ; Point back to first digit\n");
    APPEND_ESTR(helpers, "  ; Now print the string\n");
    APPEND_ESTR(helpers, "  mov rax, 1           ; sys_write\n");
    APPEND_ESTR(helpers, "  mov rsi, rdi         ; Buffer\n");
    APPEND_ESTR(helpers, "  mov rdx, rcx         ; Length\n");
    APPEND_ESTR(helpers, "  mov rdi, 1           ; stdout\n");
    APPEND_ESTR(helpers, "  syscall\n");
    APPEND_ESTR(helpers, "  mov rsp, rbp\n");
    APPEND_ESTR(helpers, "  pop rbp\n");
    APPEND_ESTR(helpers, "  ret\n");

    APPEND_ESTR(helpers, "\n; Helper: Print string at address in rax, length in rbx\n");
    APPEND_ESTR(helpers, "print_string:\n");
    APPEND_ESTR(helpers, "  push rax\n");
    APPEND_ESTR(helpers, "  push rbx\n");
    APPEND_ESTR(helpers, "  mov rdi, 1           ; stdout\n");
    APPEND_ESTR(helpers, "  mov rsi, rax         ; String address\n");
    APPEND_ESTR(helpers, "  mov rdx, rbx         ; Length\n");
    APPEND_ESTR(helpers, "  mov rax, 1           ; sys_write\n");
    APPEND_ESTR(helpers, "  syscall\n");
    APPEND_ESTR(helpers, "  pop rbx\n");
    APPEND_ESTR(helpers, "  pop rax\n");
    APPEND_ESTR(helpers, "  ret\n");

    VariableTable varTable = createVariableTable();

    // Allocate a spot for all direct references
    for (size_t i = 0; i < program->size; i++) {
        GroundInstruction gi = program->instructions[i];
        for (size_t j = 0; j < gi.args.length; j++) {
            if (gi.args.args[j].type == DIRREF) {
                addVtVariable(&varTable, gi.args.args[j].value.refName);
            }
        }
    }

    // Create data section
    for (size_t i = 0; i < varTable.count; i++) {
        if (strcmp(varTable.vars[i].name, "") == 0) {
            continue;
        }
        APPEND_ESTR(data, "  ");
        APPEND_ESTR(data, varTable.vars[i].name);
        APPEND_ESTR(data, " resq 1\n");
    }

    // Generate assembly
    for (size_t i = 0; i < program->size; i++) {
        GroundInstruction gi = program->instructions[i];
        switch (gi.type) {
            case CREATELABEL: {
                APPEND_ESTR(start, gi.args.args[0].value.refName);
                APPEND_ESTR(start, ":\n");
                break;
            }
            case SET: {
                if (gi.args.length < 2) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args", &gi, i);
                }
                if (gi.args.length > 2) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args", &gi, i);
                }
                if (gi.args.args[0].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a Value for arg 2", &gi, i);
                }

                char* varName= gi.args.args[0].value.refName;

                APPEND_ESTR(start, "  ; set\n")
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");

                break;
            }
            case JUMP: {
                if (gi.args.length < 1) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 1 arg", &gi, i);
                }
                if (gi.args.length > 1) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 1 arg", &gi, i);
                }
                if (gi.args.args[0].type != LINEREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a LineRef", &gi, i);
                }

                char* labelName = gi.args.args[0].value.refName;

                APPEND_ESTR(start, "  ; jump\n");
                APPEND_ESTR(start, "  jmp ");
                APPEND_ESTR(start, labelName);
                APPEND_ESTR(start, "\n");

                break;
            }
            case IF: {
                if (gi.args.length < 2) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args", &gi, i);
                }
                if (gi.args.length > 2) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int (for now) for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != LINEREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a LineRef", &gi, i);
                }

                char* labelName = gi.args.args[1].value.refName;

                APPEND_ESTR(start, "  ; if\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  test rax, rax\n");
                APPEND_ESTR(start, "  jnz ");
                APPEND_ESTR(start, labelName);
                APPEND_ESTR(start, "\n");

                break;
            }
            case ADD: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for add instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for add instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; add\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  add rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case SUBTRACT: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for subtract instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for subtract instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; subtract\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  sub rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case MULTIPLY: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for multiply instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for multiply instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; multiply\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  imul rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case EQUAL: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for equal instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for equal instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; equal\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  cmp rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  sete al\n");
                APPEND_ESTR(start, "  movzx rax, al\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case INEQUAL: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for inequal instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for inequal instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; inequal\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  cmp rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  setne al\n");
                APPEND_ESTR(start, "  movzx rax, al\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case GREATER: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for greater instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for greater instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; greater\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  cmp rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  setg al\n");
                APPEND_ESTR(start, "  movzx rax, al\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case LESSER: {
                if (gi.args.length < 3) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 2 args for lesser instruction", &gi, i);
                }
                if (gi.args.length > 3) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 2 args for lesser instruction", &gi, i);
                }
                if (gi.args.args[0].type != VALUE && gi.args.args[0].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 1", &gi, i);
                }
                if (gi.args.args[1].type != VALUE && gi.args.args[1].type != VALREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting an Int for arg 2", &gi, i);
                }
                if (gi.args.args[2].type != DIRREF) {
                    runtimeError(ARG_TYPE_MISMATCH, "Expecting a DirectRef for arg 3", &gi, i);
                }

                char* varName= gi.args.args[2].value.refName;
                APPEND_ESTR(start, "  ; lesser\n");
                APPEND_ESTR(start, "  mov rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  cmp rax, ");
                APPEND_ESTR(start, processValueString(gi.args.args[1]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  setl al\n");
                APPEND_ESTR(start, "  movzx rax, al\n");
                APPEND_ESTR(start, "  mov [");
                APPEND_ESTR(start, varName);
                APPEND_ESTR(start, "], rax\n");
                break;
            }
            case PRINT:
            case PRINTLN: {
                if (gi.args.length < 1) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 1 or more args", &gi, i);
                }
                
                for (size_t j = 0; j < gi.args.length; j++) {
                    if (gi.args.args[j].type != VALUE && gi.args.args[j].type != VALREF) {
                        runtimeError(ARG_TYPE_MISMATCH, "Expecting a Value", &gi, i);
                    }
                    if (j > 0) {
                        // Print space between arguments
                        APPEND_ESTR(start, "  ; print space\n");
                        APPEND_ESTR(start, "  mov rax, 1\n");
                        APPEND_ESTR(start, "  mov rdi, 1\n");
                        APPEND_ESTR(start, "  lea rsi, [rel space_char]\n");
                        APPEND_ESTR(start, "  mov rdx, 1\n");
                        APPEND_ESTR(start, "  syscall\n");
                    }
                    APPEND_ESTR(start, "  ; print int\n");
                    APPEND_ESTR(start, "  mov rax, ");
                    APPEND_ESTR(start, processValueString(gi.args.args[j]));
                    APPEND_ESTR(start, "\n");
                    APPEND_ESTR(start, "  call print_int\n");
                }
                
                if (gi.type == PRINTLN) {
                    APPEND_ESTR(start, "  ; print newline\n");
                    APPEND_ESTR(start, "  mov rax, 1\n");
                    APPEND_ESTR(start, "  mov rdi, 1\n");
                    APPEND_ESTR(start, "  lea rsi, [rel newline_char]\n");
                    APPEND_ESTR(start, "  mov rdx, 1\n");
                    APPEND_ESTR(start, "  syscall\n");
                }
                break;
            }
            case END: {
                if (gi.args.length < 1) {
                    runtimeError(TOO_FEW_ARGS, "Expecting 1 arg for end instruction", &gi, i);
                }
                if (gi.args.length > 1) {
                    runtimeError(TOO_MANY_ARGS, "Expecting 1 arg for end instruction", &gi, i);
                }
                APPEND_ESTR(start, "  ; end\n");
                APPEND_ESTR(start, "  mov rax, 60\n");
                APPEND_ESTR(start, "  mov rdi, ");
                APPEND_ESTR(start, processValueString(gi.args.args[0]));
                APPEND_ESTR(start, "\n");
                APPEND_ESTR(start, "  syscall\n");
                break;
            }
            case DROP: {
                // Drop does nothing in ground->asm as we use the stack
                break;
            }
            default: {
                printf("no\n");
                exit(1);
            }
        }
    }

    Estr complete = CREATE_ESTR("");
    APPEND_ESTR(complete, start.str);
    APPEND_ESTR(complete, "  ; End of program\n  mov rax, 60\n  mov rdi, 0\n  syscall\n");
    APPEND_ESTR(complete, helpers.str);
    APPEND_ESTR(complete, "\nsection .rodata\n");
    APPEND_ESTR(complete, "newline_char: db 10\n");
    APPEND_ESTR(complete, "space_char: db 32\n");
    APPEND_ESTR(complete, data.str)

    return complete.str;
}