#include "badlib.h"
#include "parser.c"

typedef struct Instruction {
    u8 dst;
    u8 a;
    u8 b;
    u8 op;
} Instruction;

typedef union Constant {
    s64 i;
    u64 u;
    double f;
    ptrsize ptr;
} Constant;

typedef struct Chunk {
    Instruction *code;
    Constant *constants;
    Str file_name;
    sz reg_idx;
    sz const_idx;
    Arena *storage;
    // TODO: line/col info for debugging.
} Chunk;

typedef enum OpCode {
    //            OP    DST   A   B
    // ---------------------------------------------------------------
    OP_HALT,   // halt
    OP_LD8K,   // ld8k   rx, ca      -> u8  rx = ca
    OP_LD16K,  // ld16k  rx, ca      -> u16 rx = ca
    OP_LD32K,  // ld32k  rx, ca      -> u32 rx = ca
    OP_LD64K,  // ld64k  rx, ca      -> u64 rx = ca
    OP_LD8I,   // ld8i   rx, ra, cb  -> u8  *p; rx = p[ra + cb]
    OP_LD16I,  // ld16i  rx, ra, cb  -> u16 *p; rx = p[ra + cb]
    OP_LD32I,  // ld32i  rx, ra, cb  -> u32 *p; rx = p[ra + cb]
    OP_LD64I,  // ld64i  rx, ra, cb  -> u64 *p; rx = p[ra + cb]
    OP_LD8,    // ld8    rx, ra, rb  -> u8  *p; rx = p[ra + rb]
    OP_LD16,   // ld16   rx, ra, rb  -> u16 *p; rx = p[ra + rb]
    OP_LD32,   // ld32   rx, ra, rb  -> u32 *p; rx = p[ra + rb]
    OP_LD64,   // ld64   rx, ra, rb  -> u64 *p; rx = p[ra + rb]
    OP_ST8I,   // st8i   rx, ra, cb  -> u8  *p; p[ra + cb] = rx
    OP_ST16I,  // st16i  rx, ra, cb  -> u16 *p; p[ra + cb] = rx
    OP_ST32I,  // st32i  rx, ra, cb  -> u32 *p; p[ra + cb] = rx
    OP_ST64I,  // st64i  rx, ra, cb  -> u64 *p; p[ra + cb] = rx
    OP_ST8,    // st8    rx, ra, rb  -> u8  *p; p[ra + rb] = rx
    OP_ST16,   // st16   rx, ra, rb  -> u16 *p; p[ra + rb] = rx
    OP_ST32,   // st32   rx, ra, rb  -> u32 *p; p[ra + rb] = rx
    OP_ST64,   // st64   rx, ra, rb  -> u64 *p; p[ra + rb] = rx
    OP_ADDI,   // addk   rx, ra, cb
    OP_SUBI,   // subk   rx, ra, cb
    OP_MULI,   // mulk   rx, ra, cb
    OP_DIVI,   // divk   rx, ra, cb
    OP_MODI,   // modk   rx, ra, cb
    OP_ADD,    // add    rx, ra, rb
    OP_SUB,    // sub    rx, ra, rb
    OP_MUL,    // mul    rx, ra, rb
    OP_DIV,    // div    rx, ra, rb
    OP_MOD,    // mod    rx, ra, rb
    OP_ADDFI,  // addk   rx, ra, cb
    OP_SUBFI,  // subk   rx, ra, cb
    OP_MULFI,  // mulk   rx, ra, cb
    OP_DIVFI,  // divk   rx, ra, cb
    OP_MODFI,  // modk   rx, ra, cb
    OP_ADDF,   // add    rx, ra, rb
    OP_SUBF,   // sub    rx, ra, rb
    OP_MULF,   // mul    rx, ra, rb
    OP_DIVF,   // div    rx, ra, rb
    OP_MODF,   // mod    rx, ra, rb
    OP_MOV8,   // mov8   rx, ra  -> rx = ra & 0xFF
    OP_MOV16,  // mov16  rx, ra  -> rx = ra & 0xFFFF
    OP_MOV32,  // mov32  rx, ra  -> rx = ra & 0xFFFFFFFF
    OP_MOV64,  // mov64  rx, ra  -> rx = ra & 0xFFFFFFFFFFFFFFFF
} OpCode;

Str op_str[] = {
    [OP_HALT] = cstr("HALT    "),
    // Load ops.
    [OP_LD8K] = cstr("LD8K    "),
    [OP_LD16K] = cstr("LD16K   "),
    [OP_LD32K] = cstr("LD32K   "),
    [OP_LD64K] = cstr("LD64K   "),
    [OP_LD8I] = cstr("LD8I    "),
    [OP_LD16I] = cstr("LD16I   "),
    [OP_LD32I] = cstr("LD32I   "),
    [OP_LD64I] = cstr("LD64I   "),
    [OP_LD8] = cstr("LD8     "),
    [OP_LD16] = cstr("LD16    "),
    [OP_LD32] = cstr("LD32    "),
    [OP_LD64] = cstr("LD64    "),
    // Store ops.
    [OP_ST8I] = cstr("ST8I    "),
    [OP_ST16I] = cstr("ST16I   "),
    [OP_ST32I] = cstr("ST32I   "),
    [OP_ST64I] = cstr("ST64I   "),
    [OP_ST8] = cstr("ST8     "),
    [OP_ST16] = cstr("ST16    "),
    [OP_ST32] = cstr("ST32    "),
    [OP_ST64] = cstr("ST64    "),
    // Arithmetic.
    [OP_ADDI] = cstr("ADDI    "),
    [OP_SUBI] = cstr("SUBI    "),
    [OP_MULI] = cstr("MULI    "),
    [OP_DIVI] = cstr("DIVI    "),
    [OP_MODI] = cstr("MODI    "),
    [OP_ADD] = cstr("ADD     "),
    [OP_SUB] = cstr("SUB     "),
    [OP_MUL] = cstr("MUL     "),
    [OP_DIV] = cstr("DIV     "),
    [OP_MOD] = cstr("MOD     "),
    [OP_ADDFI] = cstr("ADDFI   "),
    [OP_SUBFI] = cstr("SUBFI   "),
    [OP_MULFI] = cstr("MULFI   "),
    [OP_DIVFI] = cstr("DIVFI   "),
    [OP_MODFI] = cstr("MODFI   "),
    [OP_ADDF] = cstr("ADDF    "),
    [OP_SUBF] = cstr("SUBF    "),
    [OP_MULF] = cstr("MULF    "),
    [OP_DIVF] = cstr("DIVF    "),
    // Reg copy/move.
    [OP_MODF] = cstr("MODF    "),
    [OP_MOV8] = cstr("MOV8    "),
    [OP_MOV16] = cstr("MOV16   "),
    [OP_MOV32] = cstr("MOV32   "),
    [OP_MOV64] = cstr("MOV64   "),
};

void
disassemble_instruction(Instruction instruction) {
    switch (instruction.op) {
        case OP_MOV8:
        case OP_MOV16:
        case OP_MOV32:
        case OP_MOV64:
            println("%s r%d, r%d", op_str[instruction.op], instruction.dst,
                    instruction.a, instruction.b);
            break;
        case OP_LD8K:
        case OP_LD16K:
        case OP_LD32K:
        case OP_LD64K:
            println("%s r%d, c%d", op_str[instruction.op], instruction.dst,
                    instruction.a, instruction.b);
            break;
        case OP_LD8I:
        case OP_LD16I:
        case OP_LD32I:
        case OP_LD64I:
        case OP_ST8I:
        case OP_ST16I:
        case OP_ST32I:
        case OP_ST64I:
        case OP_ADDI:
        case OP_SUBI:
        case OP_MULI:
        case OP_DIVI:
        case OP_MODI:
        case OP_ADDFI:
        case OP_SUBFI:
        case OP_MULFI:
        case OP_DIVFI:
        case OP_MODFI:
            println("%s r%d, r%d, c%d", op_str[instruction.op], instruction.dst,
                    instruction.a, instruction.b);
            break;
        case OP_LD8:
        case OP_LD16:
        case OP_LD32:
        case OP_LD64:
        case OP_ST8:
        case OP_ST16:
        case OP_ST32:
        case OP_ST64:
        case OP_ADD:
        case OP_SUB:
        case OP_MUL:
        case OP_DIV:
        case OP_MOD:
        case OP_ADDF:
        case OP_SUBF:
        case OP_MULF:
        case OP_DIVF:
        case OP_MODF:
            println("%s r%d, r%d, r%d", op_str[instruction.op], instruction.dst,
                    instruction.a, instruction.b);
            break;
        case OP_HALT: println("%s", op_str[instruction.op]); break;
        default: println("Unknown opcode %d", instruction.op); break;
    }
}

void
disassemble_chunk(Chunk chunk) {
    println("%s: =========== code ===========", chunk.file_name);
    for (sz i = 0; i < array_size(chunk.code); i++) {
        print("%s: %x{4}:  ", chunk.file_name, i);
        disassemble_instruction(chunk.code[i]);
    }
    if (array_size(chunk.constants) > 0) {
        println("%s: ========= constants ========", chunk.file_name);
        for (sz i = 0; i < array_size(chunk.constants); i++) {
            println("%s: %x{2}:  %x{8}", chunk.file_name, i,
                    chunk.constants[i]);
        }
    }
}

#define N_CONST 256
typedef struct VM {
    Chunk *chunk;
    Constant regs[N_CONST];
    Instruction *ip;
} VM;

void
vm_init(VM *vm, Chunk *chunk) {
    assert(vm);
    assert(chunk);
    assert(chunk->code);
    vm->chunk = chunk;
    vm->ip = vm->chunk->code;
}

#define OP_ARITHMETIC(OP, TYPE)                                            \
    do {                                                                   \
        u8 dst = instruction.dst;                                          \
        u8 src_a = instruction.a;                                          \
        u8 src_b = instruction.b;                                          \
        vm->regs[dst].TYPE = vm->regs[src_a].TYPE OP vm->regs[src_b].TYPE; \
    } while (0);

#define OP_ARITHMETIC_CONST(OP, TYPE)                                 \
    do {                                                              \
        u8 dst = instruction.dst;                                     \
        u8 src_a = instruction.a;                                     \
        u8 src_b = instruction.b;                                     \
        vm->regs[dst].TYPE =                                          \
            vm->regs[src_a].TYPE OP vm->chunk->constants[src_b].TYPE; \
    } while (0);

#include <math.h>

void
vm_run(VM *vm) {
    assert(vm);
    assert(vm->chunk);
    assert(vm->ip);
    println("VM running...");
    while (true) {
        Instruction instruction = *vm->ip++;
#if DEBUG == 1
        print("IP: %d  -> ", vm->ip - vm->chunk->code - 1);
        disassemble_instruction(instruction);
#endif

        switch (instruction.op) {
            case OP_LD64K: {
                u8 dst = instruction.dst;
                u8 src_a = instruction.a;
                vm->regs[dst].i = vm->chunk->constants[src_a].i;
            } break;
            case OP_ADD: OP_ARITHMETIC(+, i) break;
            case OP_SUB: OP_ARITHMETIC(-, i) break;
            case OP_MUL: OP_ARITHMETIC(*, i) break;
            case OP_DIV: OP_ARITHMETIC(/, i) break;
            case OP_MOD: OP_ARITHMETIC(%, i) break;
            case OP_ADDF: OP_ARITHMETIC(+, f) break;
            case OP_SUBF: OP_ARITHMETIC(-, f) break;
            case OP_MULF: OP_ARITHMETIC(*, f) break;
            case OP_DIVF: OP_ARITHMETIC(/, f) break;
            case OP_MODF: {
                u8 dst = instruction.dst;
                u8 src_a = instruction.a;
                u8 src_b = instruction.b;
                vm->regs[dst].f =
                    fmod(vm->regs[src_a].f, vm->chunk->constants[src_b].f);
            } break;
            case OP_HALT: {
                println("VM HALT -> %d", vm->regs[instruction.dst]);
                return;
            }

            default: {
                eprintln("unimplemented OP code: %d", instruction.op);
                return;
            }
        }
    }
}

typedef enum {
    COMP_CONST,
    COMP_REG,
    COMP_ERR,
} CompResultType;

typedef struct CompResult {
    sz idx;
    CompResultType type;
} CompResult;

CompResult compile_expr(Chunk *chunk, Node *node);

// #define EMIT_OP(OP, CHUNK, ARENA)

CompResult
compile_binary(OpCode op, Chunk *chunk, Node *node) {
    sz reg_dst = chunk->reg_idx++;
    CompResult comp_a = compile_expr(chunk, node->left);
    CompResult comp_b = compile_expr(chunk, node->right);
    sz reg_a;
    sz reg_b;
    switch (comp_a.type) {
        case COMP_CONST: {
            reg_a = chunk->reg_idx++;
            Instruction inst =
                (Instruction){.op = OP_LD64K, .dst = reg_a, .a = comp_a.idx};
            array_push(chunk->code, inst, chunk->storage);
        } break;
        case COMP_REG: {
            reg_a = comp_a.idx;
        } break;
        default: {
            return (CompResult){.type = COMP_ERR};
        } break;
    }
    switch (comp_b.type) {
        case COMP_CONST: {
            reg_b = chunk->reg_idx++;
            Instruction inst =
                (Instruction){.op = OP_LD64K, .dst = reg_b, .a = comp_b.idx};
            array_push(chunk->code, inst, chunk->storage);
        } break;
        case COMP_REG: {
            reg_b = comp_b.idx;
        } break;
        default: {
            return (CompResult){.type = COMP_ERR};
        } break;
    }
    Instruction inst =
        (Instruction){.op = op, .dst = reg_dst, .a = reg_a, .b = reg_b};
    array_push(chunk->code, inst, chunk->storage);
    return (CompResult){.type = COMP_REG, .idx = reg_dst};
}

CompResult
compile_expr(Chunk *chunk, Node *node) {
    switch (node->kind) {
        case NODE_ADD: return compile_binary(OP_ADD, chunk, node); break;
        case NODE_SUB: return compile_binary(OP_SUB, chunk, node); break;
        case NODE_MUL: return compile_binary(OP_MUL, chunk, node); break;
        case NODE_DIV: return compile_binary(OP_DIV, chunk, node); break;
        case NODE_MOD: return compile_binary(OP_MOD, chunk, node); break;
        case NODE_NUM_FLOAT:
        case NODE_NUM_INT: {
            // Make sure we don't have duplicated constants.
            for (sz i = 0; i < chunk->const_idx; i++) {
                if (node->value.i == chunk->constants[i].i) {
                    return (CompResult){
                        .type = COMP_CONST,
                        .idx = i,
                    };
                }
            }
            Constant c = (Constant){.i = node->value.i};
            array_push(chunk->constants, c, chunk->storage);
            return (CompResult){
                .type = COMP_CONST,
                .idx = chunk->const_idx++,
            };
        } break;
        default: break;
    }
    return (CompResult){.type = COMP_ERR};
}