#include "argparse.h"
#include "dynarray.h"
#include "logger.h"
#include "parser/ast.h"
#include "token.h"
#include <stdio.h>
#include <string.h>

typedef struct comp_state_s {
    char** strings;
    size_t if_id;
    size_t while_id;
} comp_state_t;

int write_if_stat(ast_if_stat_t* if_stat, FILE* f, comp_state_t* state);

int write_op(ast_op_t* aop, FILE* f, comp_state_t* state) {

    switch (aop->type) {
    case (AOT_OP): {
        token_t op = aop->op;
        switch (op.type) {
        case (TT_PUSH_INT): {
            fprintf(f, "    ; -- PUSH_INT --\n");
            fprintf(f, "    mov rax, %zu\n", op.int_v);
            fprintf(f, "    push rax\n");
        }; break;
        case (TT_PUSH_STR): {
            fprintf(f, "    ; -- PUSH_STR --\n");
            fprintf(f, "    mov rax, morph_str_%zu\n", dynarray_length(state->strings));
            fprintf(f, "    push rax\n");
            fprintf(f, "    mov rax, %zu\n", strlen(op.str_v));
            fprintf(f, "    push rax\n");
            dynarray_push(state->strings, op.str_v);
        }; break;
        case (TT_PUSH_CSTR): {
            fprintf(f, "    ; -- PUSH_CSTR --\n");
            fprintf(f, "    mov rax, morph_str_%zu\n", dynarray_length(state->strings));
            fprintf(f, "    push rax\n");
            dynarray_push(state->strings, op.str_v);
        }; break;
        case (TT_PUSH_BOOL): {
            fprintf(f, "    ; -- PUSH_BOOL --\n");
            fprintf(f, "    mov rax, %d\n", (int)op.bool_v);
            fprintf(f, "    push rax\n");
        }; break;
        case (TT_PUSH_CHAR): {
            fprintf(f, "    ; -- PUSH_CHAR --\n");
            fprintf(f, "    mov rax, %d ; '%c'\n", (int)op.char_v, op.char_v);
            fprintf(f, "    push rax\n");
        }; break;
        case (TT_PUSH_FLOAT): {
            fprintf(f, "    ; -- PUSH_FLOAT --\n");
            // TODO: do this
        }; break;
        case (TT_OP): {
            switch (op.op_type) {

            case (OP_CAST_BOOL):
            case (OP_CAST_INT):
            case (OP_CAST_PTR):
            case (OP_COUNT__):
            case (OP_NONE):
                break;
            case (OP_ADD): {
                fprintf(f, "    ; -- OP_ADD --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    add rax, rbx\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SUB): {
                fprintf(f, "    ; -- OP_SUB --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    sub rax, rbx\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_MUL): {
                fprintf(f, "    ; -- OP_MUL --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    mul rbx\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_DIV): {
                fprintf(f, "    ; -- OP_DIV --\n");
                fprintf(f, "    xor rdx, rdx\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    div rbx\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_MOD): {
                fprintf(f, "    ; -- OP_MOD --\n");
                fprintf(f, "    xor rdx, rdx\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    div rbx\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_EQ): {
                fprintf(f, "    ; -- OP_EQ --\n");
                fprintf(f, "    mov rcx, 0\n");
                fprintf(f, "    mov rdx, 1\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmove rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_GT): {
                fprintf(f, "    ; -- OP_GT --\n");
                fprintf(f, "    mov rcx, 0\n");
                fprintf(f, "    mov rdx, 1\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmovg rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_LT): {
                fprintf(f, "    ; -- OP_LT --\n");
                fprintf(f, "    mov rcx, 0\n");
                fprintf(f, "    mov rdx, 1\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmovl rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_GE): {
                fprintf(f, "    ; -- OP_GE --\n");
                fprintf(f, "    mov rcx, 0\n");
                fprintf(f, "    mov rdx, 1\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmovge rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_LE): {
                fprintf(f, "    ; -- OP_LE --\n");
                fprintf(f, "    mov rcx, 0\n");
                fprintf(f, "    mov rdx, 1\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmovle rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_NE): {
                fprintf(f, "    ; -- OP_NE --\n");
                fprintf(f, "    mov rcx, 1\n");
                fprintf(f, "    mov rdx, 0\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    cmp rax, rbx\n");
                fprintf(f, "    cmove rcx, rdx\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_AND): {
                fprintf(f, "    ; -- OP_AND --\n");
                fprintf(f, "    pop     rax\n");
                fprintf(f, "    pop     rbx\n");
                fprintf(f, "    test    rax, rax\n");
                fprintf(f, "    setnz   al\n");
                fprintf(f, "    test    rbx, rbx\n");
                fprintf(f, "    setnz   bl\n");
                fprintf(f, "    and     al, bl\n");
                fprintf(f, "    movzx   rbx, al\n");
                fprintf(f, "    push    rbx\n");
            }; break;
            case (OP_OR): {
                fprintf(f, "    ; -- OP_OR --\n");
                fprintf(f, "    pop     rax\n");
                fprintf(f, "    pop     rbx\n");
                fprintf(f, "    test    rax, rax\n");
                fprintf(f, "    setnz   al\n");
                fprintf(f, "    test    rbx, rbx\n");
                fprintf(f, "    setnz   bl\n");
                fprintf(f, "    or      al, bl\n");
                fprintf(f, "    movzx   rbx, al\n");
                fprintf(f, "    push    rbx\n");
            }; break;
            case (OP_SHR): {
                fprintf(f, "    ; -- OP_SHR --\n");
                fprintf(f, "    pop rcx\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    shr rbx, cl\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_SHL): {
                fprintf(f, "    ; -- OP_SHL --\n");
                fprintf(f, "    pop rcx\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    shl rbx, cl\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_BOR): {
                fprintf(f, "    ; -- OP_BOR --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    or rbx, rax\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_BAND): {
                fprintf(f, "    ; -- OP_BAND --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    and rbx, rax\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_NOT): {
                fprintf(f, "    ; -- OP_NOT --\n");
                fprintf(f, "    pop     rax\n");
                fprintf(f, "    test    rax, rax\n");
                fprintf(f, "    setz    al\n");
                fprintf(f, "    movzx   rbx, al\n");
                fprintf(f, "    push    rbx\n");
            }; break;
            case (OP_DUP): {
                fprintf(f, "    ; -- OP_DUP --\n");
                fprintf(f, "    pop      rax\n");
                fprintf(f, "    push     rax\n");
                fprintf(f, "    push     rax\n");
            }; break;
            case (OP_SWAP): {
                fprintf(f, "    ; -- OP_SWAP --\n");
                fprintf(f, "    pop      rax\n");
                fprintf(f, "    pop      rbx\n");
                fprintf(f, "    push     rax\n");
                fprintf(f, "    push     rbx\n");
            }; break;
            case (OP_DROP): {
                fprintf(f, "    ; -- OP_DROP --\n");
                fprintf(f, "    pop      rax\n");
            }; break;
            case (OP_OVER): {
                fprintf(f, "    ; -- OP_OVER --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    push rbx\n");
                fprintf(f, "    push rax\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_ROT): {
                fprintf(f, "    ; -- OP_ROT --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rcx\n");
                fprintf(f, "    push rbx\n");
                fprintf(f, "    push rax\n");
                fprintf(f, "    push rcx\n");
            }; break;
            case (OP_LOAD8): {
                fprintf(f, "    ; -- OP_LOAD8 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    xor rbx, rbx\n");
                fprintf(f, "    mov bl, byte [rax]\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_STORE8): {
                fprintf(f, "    ; -- OP_STORE8 --\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    mov byte [rax], bl\n");
            }; break;
            case (OP_LOAD16): {
                fprintf(f, "    ; -- OP_LOAD16 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    xor rbx, rbx\n");
                fprintf(f, "    mov bx, word [rax]\n");
                fprintf(f, "    push rbx\n");
            }; break;

            case (OP_STORE16): {
                fprintf(f, "    ; -- OP_STORE16 --\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    mov word [rax], bx\n");
            }; break;
            case (OP_LOAD32): {
                fprintf(f, "    ; -- OP_LOAD32 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    xor rbx, rbx\n");
                fprintf(f, "    mov ebx, dword [rax]\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_STORE32): {
                fprintf(f, "    ; -- OP_STORE32 --\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    mov dword[rax], ebx\n");
            }; break;
            case (OP_LOAD64): {
                fprintf(f, "    ; -- OP_LOAD64 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    xor rbx, rbx\n");
                fprintf(f, "    mov rbx, qword [rax]\n");
                fprintf(f, "    push rbx\n");
            }; break;
            case (OP_STORE64): {
                fprintf(f, "    ; -- OP_STORE64 --\n");
                fprintf(f, "    pop rbx\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    mov qword [rax], rbx\n");
            }; break;
            case (OP_SYSCALL0): {
                fprintf(f, "    ; -- OP_SYSCALL0 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL1): {
                fprintf(f, "    ; -- OP_SYSCALL1 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop r9\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL2): {
                fprintf(f, "    ; -- OP_SYSCALL2 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop rsi\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL3): {
                fprintf(f, "    ; -- OP_SYSCALL3 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop rsi\n");
                fprintf(f, "    pop rdx\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL4): {
                fprintf(f, "    ; -- OP_SYSCALL4 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop rsi\n");
                fprintf(f, "    pop rdx\n");
                fprintf(f, "    pop r10\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL5): {
                fprintf(f, "    ; -- OP_SYSCALL5 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop rsi\n");
                fprintf(f, "    pop rdx\n");
                fprintf(f, "    pop r10\n");
                fprintf(f, "    pop r8\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_SYSCALL6): {
                fprintf(f, "    ; -- OP_SYSCALL6 --\n");
                fprintf(f, "    pop rax\n");
                fprintf(f, "    pop rdi\n");
                fprintf(f, "    pop rsi\n");
                fprintf(f, "    pop rdx\n");
                fprintf(f, "    pop r10\n");
                fprintf(f, "    pop r8\n");
                fprintf(f, "    pop r9\n");
                fprintf(f, "    syscall\n");
                fprintf(f, "    push rax\n");
            }; break;
            case (OP_ARGC): {
            }; break;
            case (OP_ARGV): {
            }; break;
            case (OP_HERE): {
                char* pos = malloc(1024 * 2);
                snprintf(pos, 2048, "%s:%d:%d", op.loc.file, op.loc.line, op.loc.col);
                fprintf(f, "    mov rax, str_%zu\n", dynarray_length(state->strings));
                fprintf(f, "    push rax\n");
                fprintf(f, "    mov rax, %zu\n", strlen(op.str_v));
                fprintf(f, "    push rax\n");
                dynarray_push(state->strings, op.str_v);
            }; break;
            case (OP_PRINT): {
            }; break;
            }
        }; break;
        case (TT_KW):
        case (TT_IDENT): {
            log_warn(&op.loc, "Found a KW or IDENT where it shouldnt exists (compiler): %s", get_tok_str_dbg(&op));
            // unreachable
        } break;
        }
    }
    case (AOT_IF): {
        write_if_stat(&aop->if_stat, f, state);
    }; break;
    case (AOT_WHILE): {
        size_t id = state->while_id++;
        fprintf(f, " ;; -- OP_WHILE -- \n");
        fprintf(f, "morph_while_%zu_cond:\n", id);
        for (int i = 0; i < dynarray_length(aop->while_stat.condition); i++) {
            write_op(&aop->while_stat.condition[i], f, state);
        }
        fprintf(f, "    pop rax\n");
        fprintf(f, "    test rax, rax\n");
        fprintf(f, "    jz morph_while_%zu_end\n", id);
        fprintf(f, "    pop rax\n");
        fprintf(f, "morph_while_%zu_start:\n", id);
        for (int i = 0; i < dynarray_length(aop->while_stat.body); i++) {
            write_op(&aop->while_stat.body[i], f, state);
        }
        fprintf(f, "    jmp morph_while_%zu_cond:\n", id);
        fprintf(f, "morph_while_%zu_end:\n", id);
    }; break;
    case (AOT_USE_CONST): {

        fprintf(f, "    ;; -- OP_USE_CONST -- \n");
        fprintf(f, "    mov rax, [morph_const_%zu]\n", aop->id);
        fprintf(f, "    push rax\n");
    }; break;
    case (AOT_USE_MEMORY): {
        fprintf(f, "    ;; -- OP_USE_MEMORY -- \n");
        fprintf(f, "    mov rax, morph_memory_%zu\n", aop->id);
        fprintf(f, "    push rax\n");
    }; break;
    case (AOT_CALL_FUNC): {

    }; break;
    }
    return 0;
}

int write_if_stat(ast_if_stat_t* if_stat, FILE* f, comp_state_t* state) {

    size_t id = state->if_id++;
    fprintf(f, "    ; -- OP_IF \n");
    for (int i = 0; i < dynarray_length(if_stat->condition); i++) {
        write_op(&if_stat->condition[i], f, state);
    }

    fprintf(f, "    pop rax\n");
    fprintf(f, "    test rax, rax\n");
    fprintf(f, "    jz morph_if_%zu_cond_false\n", id);
    fprintf(f, "    pop rax\n");
    fprintf(f, "morph_if_%zu_cond_true:\n", id);
    for (int i = 0; i < dynarray_length(if_stat->body); i++) {
        write_op(&if_stat->body[i], f, state);
    }
    fprintf(f, "    jmp morph_if_%zu_end\n", id);
    fprintf(f, "morph_if_%zu_cond_false:\n", id);
    if (if_stat->else_body) {
        for (int i = 0; i < dynarray_length(if_stat->body); i++) {
            write_op(&if_stat->body[i], f, state);
        }
    } else if (if_stat->is_elseif) {
        write_if_stat(if_stat->elseif, f, state);
    }
    fprintf(f, "morph_if_%zu_end:\n", id);
    return 0;
}

int compile_x86_64_linux_nasm(args_t* args, program_t* prog) {
    FILE* f = fopen(args->asm_file, "w");
    comp_state_t state = {0};
    state.strings = dynarray_create(char*);

    fprintf(f, "bits 64\n");
    fprintf(f, "section .text\n");
    fprintf(f, "global _start\n");
    fprintf(f, "_start:\n");
    fprintf(f, "    mov rax, [rsp]\n");
    fprintf(f, "    mov [morph_i_argc], rax\n");
    fprintf(f, "    mov rax, [rsp+8]\n");
    fprintf(f, "    mov [morph_i_argv], rax\n");
    fprintf(f, "    call morph_f_main\n");
    fprintf(f, "    mov rax, 60\n");
    fprintf(f, "    mov rdi, 0\n");
    fprintf(f, "    syscall\n");
    size_t fn_len = dynarray_length(prog->funcs);
    for (int i = 0; i < fn_len; i++) {
        function_t func = prog->funcs[i];

        fprintf(f, "morph_f_%s:\n", func.name);
        for (int y = 0; y < dynarray_length(func.body); y++) {
            ast_op_t aop = func.body[y];
            if (write_op(&aop, f, &state)) {
                return 1;
            }
        }
        fprintf(f, "    ret\n");
    }

    fprintf(f, "section .bss\n");
    fprintf(f, "morph_i_argc: resb 8\n");
    fprintf(f, "morph_i_argv: resb 8\n");
    for (int i = 0; i < dynarray_length(prog->memories); i++) {
        memory_t mem = prog->memories[i];
        fprintf(f, "morph_memory_%d: resb %zu ; Memory %s\n", i, mem.size, mem.name);
    }
    fprintf(f, "section .rodata\n");

    for (int i = 0; i < dynarray_length(state.strings); i++) {
        fprintf(f, "morph_str_%d: db \"%s\"\n", i, state.strings[i]);
    }
    for (int i = 0; i < dynarray_length(prog->const_vars); i++) {
        const_t* v = &prog->const_vars[i];
        fprintf(f, "morph_const_%d: ", i);
        switch (v->val.type) {
        case (TT_PUSH_INT): {
            fprintf(f, "dq %zu\n", v->val.int_v);
        }; break;
        case (TT_PUSH_CSTR):
        case (TT_PUSH_STR): {
            fprintf(f, "db \"%s\"\n", v->val.str_v);
        }; break;
        case (TT_PUSH_CHAR):
            fprintf(f, "db %c\n", v->val.char_v);
        case (TT_PUSH_BOOL):
            fprintf(f, "db %d\n", (int)v->val.bool_v);
        default:
        }
    }
    return 0;
}