#ifndef BDL_VM_H
#define BDL_VM_H

#include "types.h"
#include "errors.h"
#include "chunk.h"
#include "ops.h"
#include "debug.h"
#include "hashtable.h"

#define VM_FRAMES_CAP 64
#define VM_STACK_CAP 1024

typedef struct CallFrame {
    // Current code being run.
    Chunk *chunk;
    // Current program counter for this call.
    // Ref to stack. Is this really needed?
    // Object *locals;
    // Return counter.
    u8 *rp;
    // Starting point of the locals for this procedure.
    size_t stack_offset;
} CallFrame;

typedef struct VM {
    CallFrame *frames;
    // Stack.
    Object *stack;
    // Program counter.
    u8 *pc;
    // Global variables.
    HashTable *globals;
} VM;

void vm_init(VM *vm);
void vm_free(VM *vm);
void vm_reset(VM *vm);
void vm_interpret(VM *vm);

void
vm_init(VM *vm) {
    *vm = (VM){0};
    array_init(vm->frames, VM_FRAMES_CAP);
    array_init(vm->stack, VM_STACK_CAP);
    vm->globals = ht_init();
}

void
vm_free(VM *vm) {
    array_free(vm->frames);
    array_free(vm->stack);
    ht_free(vm->globals);
}

void
vm_reset(VM *vm) {
    vm_free(vm);
    vm_init(vm);
}

// Helper macros for a more clear VM switch.
#ifdef DEBUG
#define STACK_TRACE() \
    fprintf(stderr, "stack trace:\n"); \
    for (ssize_t i = array_size(vm->frames) - 1; i >= 0 ; i--) { \
        CallFrame frame = vm->frames[i]; \
        Chunk *chunk = frame.chunk; \
        size_t instruction = vm->pc - chunk->code - 1; \
        fprintf(stderr, "\t%-4ld  ->  ", i); \
        fprintf(stderr, "%.*s",(int)array_size(chunk->name), chunk->name); \
        fprintf(stderr, ":%ld:%ld\n", chunk->lines[instruction].line, chunk->lines[instruction].col); \
        vm->pc = frame.rp; \
    }
#else
#define STACK_TRACE()
#endif

#define RUNTIME_ERROR(ERR) \
    error_push((Error){ \
        .type = ERR_TYPE_RUNTIME, \
        .value = (ERR), \
        .line = frame->chunk->lines[vm->pc - frame->chunk->code - 1].line, \
        .col = frame->chunk->lines[vm->pc - frame->chunk->code - 1].col, \
    }); \
    STACK_TRACE() \
    return

#define FIXNUM_ARITHMETIC_OP(OP) \
    do { \
        ssize_t n = AS_FIXNUM(array_pop(vm->stack)); \
        size_t stack_size = array_size(vm->stack) - n; \
        Object obj = array_peek(vm->stack, n - 1); \
        if (!IS_FIXNUM(obj)) { \
            RUNTIME_ERROR(ERR_WRONG_ARG_TYPE); \
            return; \
        } \
        ssize_t acc = AS_FIXNUM(obj); \
        while (n > 1) { \
            obj = array_peek(vm->stack, n - 2); \
            if (!IS_FIXNUM(obj)) { \
                RUNTIME_ERROR(ERR_WRONG_ARG_TYPE); \
                return; \
            } \
            ssize_t current = AS_FIXNUM(obj); \
            acc = acc OP current; \
            n--; \
        } \
        array_head(vm->stack)->size = stack_size; \
        array_push(vm->stack, FIXNUM_VAL(acc)); \
    } while (false)

#define FIXNUM_COMPARE_OP(OP) \
    do { \
        ssize_t n = AS_FIXNUM(array_pop(vm->stack)); \
        size_t stack_size = array_size(vm->stack) - n; \
        Object obj = array_peek(vm->stack, n - 1); \
        if (!IS_FIXNUM(obj)) { \
            RUNTIME_ERROR(ERR_WRONG_ARG_TYPE); \
            return; \
        } \
        ssize_t prev = AS_FIXNUM(obj); \
        bool ret = true; \
        while (n > 1) { \
            obj = array_peek(vm->stack, n - 2); \
            if (!IS_FIXNUM(obj)) { \
                RUNTIME_ERROR(ERR_WRONG_ARG_TYPE); \
                return; \
            } \
            ssize_t current = AS_FIXNUM(obj); \
            ret = ret && (prev OP current); \
            prev = current; \
            n--; \
        } \
        array_head(vm->stack)->size = stack_size; \
        array_push(vm->stack, BOOL_VAL(ret)); \
    } while (false)

#define LOGIC_OP(OP) \
    do { \
        Object a = array_pop(vm->stack); \
        Object b = array_pop(vm->stack); \
        bool x = IS_TRUE(a); \
        bool y = IS_TRUE(b); \
        Object result = y OP x ? TRUE_VAL : FALSE_VAL; \
        array_push(vm->stack, result); \
    } while (false)

void
vm_interpret(VM *vm) {
    CallFrame *frame = &vm->frames[0];
    vm->pc = frame->chunk->code;
    frame->rp = NULL;

    if (frame->chunk->code == NULL || array_size(frame->chunk->code) == 0) {
        error_push((Error){
            .type = ERR_TYPE_RUNTIME,
            .value = ERR_EMPTY_CHUNK,
        });
        return;
    }

    while (true) {
#ifdef DEBUG_TRACE_EXECUTION
        printf("stack:  [ ");
        for (size_t i = 0; i < array_size(vm->stack); i++) {
            object_display(vm->stack[i]);
            if (i < array_size(vm->stack) - 1) {
                printf(" | ");
            }
        }
        printf(" ]\nop: ");
        disassemble_instruction(frame->chunk, (vm->pc - frame->chunk->code));
#endif
        u8 instruction = *vm->pc++;
        switch (instruction) {
            case OP_LOCAL: {
                u8 local = *vm->pc++;
                array_push(vm->stack, vm->stack[frame->stack_offset + local]);
            } break;
            case OP_CONSTANT: {
                u8 constant = *vm->pc++;
                Object obj = frame->chunk->constants[constant];
                array_push(vm->stack, obj);
            } break;
            case OP_DEF_LOCAL: {
                Object value = array_pop(vm->stack);
                ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
                vm->stack[frame->stack_offset + idx] = value;
            } break;
            case OP_SET_LOCAL: {
                Object value = array_pop(vm->stack);
                ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
                ssize_t depth = AS_FIXNUM(array_pop(vm->stack));
                CallFrame frame = vm->frames[depth];
                vm->stack[frame.stack_offset + idx] = value;
            } break;
            case OP_DEF: {
                Object value = array_pop(vm->stack);
                Object name = array_pop(vm->stack);
                ht_insert(vm->globals, name, value);
            } break;
            case OP_SET: {
                Object value = array_pop(vm->stack);
                Object name = array_pop(vm->stack);
                Object *prev = ht_lookup(vm->globals, name);
                if (prev == NULL) {
                    RUNTIME_ERROR(ERR_SYMBOL_NOT_FOUND);
                    return;
                }
                ht_insert(vm->globals, name, value);
            } break;
            case OP_GET: {
                Object name = array_pop(vm->stack);
                Object *value = ht_lookup(vm->globals, name);
                if (value == NULL) {
                    RUNTIME_ERROR(ERR_SYMBOL_NOT_FOUND);
                    return;
                }
                array_push(vm->stack, *value);
            } break;
            case OP_SUM: { FIXNUM_ARITHMETIC_OP(+); } break;
            case OP_SUB: { FIXNUM_ARITHMETIC_OP(-); } break;
            case OP_MUL: { FIXNUM_ARITHMETIC_OP(*); } break;
            case OP_DIV: { FIXNUM_ARITHMETIC_OP(/); } break;
            case OP_MOD: { FIXNUM_ARITHMETIC_OP(%); } break;
            case OP_NOT: {
                Object prev = array_pop(vm->stack);
                Object new = IS_TRUE(prev) ? FALSE_VAL : TRUE_VAL;
                array_push(vm->stack, new);
            } break;
            case OP_AND: { LOGIC_OP(&&); } break;
            case OP_OR: { LOGIC_OP(||); } break;
            case OP_EQUAL: { FIXNUM_COMPARE_OP(==); } break;
            case OP_LESS: { FIXNUM_COMPARE_OP(<); } break;
            case OP_GREATER: { FIXNUM_COMPARE_OP(>); } break;
            case OP_LESS_EQUAL: { FIXNUM_COMPARE_OP(<=); } break;
            case OP_GREATER_EQUAL: { FIXNUM_COMPARE_OP(>=); } break;
            case OP_JUMP: {
                u16 a = *vm->pc++;
                u16 b = *vm->pc++;
                s16 offset = (a << 8) | b;
                vm->pc += offset;
            } break;
            case OP_JUMP_IF_FALSE: {
                u16 a = *vm->pc++;
                u16 b = *vm->pc++;
                s16 offset = (a << 8) | b;
                if (IS_FALSE(array_pop(vm->stack))) {
                    vm->pc += offset;
                }
            } break;
            case OP_DISPLAY: {
                object_display(array_pop(vm->stack));
            } break;
            case OP_PRINT: {
                Object obj = array_pop(vm->stack);
                if (!IS_STRING(obj)) {
                    RUNTIME_ERROR(ERR_WRONG_ARG_TYPE);
                }
                StringView scanner = (StringView) {
                    .start = obj.text,
                    .n = array_size(obj.text),
                };
                while (scanner.n != 0) {
                    char c = sv_next(&scanner);
                    if (c == '\\' && sv_peek(&scanner) == 'n') {
                        putchar('\n');
                        sv_next(&scanner);
                        continue;
                    }
                    if (c == '\\' && sv_peek(&scanner) == '"') {
                        putchar('"');
                        sv_next(&scanner);
                        continue;
                    }
                    putchar(c);
                }
            } break;
            case OP_NEWLINE: {
                printf("\n");
            } break;
            case OP_CALL: {
                ssize_t n_args = AS_FIXNUM(array_pop(vm->stack));
                Object proc = array_pop(vm->stack);

                // Check the number of arguments is correct.
                // NOTE: This is probably better handled at compilation, but for
                // now this is simpler to implement.
                ssize_t n_params = proc.chunk->n_params;
                ssize_t n_locals = proc.chunk->n_locals;
                if (n_args < n_params) {
                    RUNTIME_ERROR(ERR_NOT_ENOUGH_ARGS);
                } else if (n_args > n_params) {
                    RUNTIME_ERROR(ERR_TOO_MANY_ARGS);
                }

#ifdef DEBUG
                disassemble_chunk(proc.chunk);
#endif

                // Tail-call optimization.
                if (proc.chunk != frame->chunk || *vm->pc != OP_RETURN) {
                    // Prepare new call frame.
                    CallFrame new_frame = (CallFrame){
                        .chunk = proc.chunk,
                        .rp = vm->pc,
                        .stack_offset = array_size(vm->stack) - n_params,
                    };
                    array_push(vm->frames, new_frame);
                    frame = &vm->frames[array_size(vm->frames) - 1];

                    // Prepare local slots.
                    for (ssize_t i = 0; i < n_locals; i++) {
                        array_push(vm->stack, NIL_VAL);
                    }
                } else {
                    // Bind tail-call parameters.
                    for (ssize_t i = 0; i < n_params; i++) {
                        Object obj = array_peek(vm->stack, n_locals + n_params - 1 - i);
                        vm->stack[frame->stack_offset + i] = obj;
                    }

                    // Reset stack size.
                    size_t offset = frame->stack_offset + n_params + n_locals;
                    array_head(vm->stack)->size = offset;
                }
                vm->pc = frame->chunk->code;
            } break;
            case OP_RETURN: {
                if (frame->rp == NULL) {
                    Object ret = array_pop(vm->stack);
                    if (!IS_NIL(ret)) {
                        object_display(ret);
                        printf("\n");
                    }
                    return;
                }
                vm->pc = frame->rp;
                array_head(vm->frames)->size--;
                Object ret = array_pop(vm->stack);
                array_head(vm->stack)->size = frame->stack_offset;
                array_push(vm->stack, ret);
                frame = &vm->frames[array_size(vm->frames) - 1];
            } break;
            case OP_DROP: {
                array_head(vm->stack)->size = 0;
            } break;
            default: {
                RUNTIME_ERROR(ERR_NOT_IMPLEMENTED);
            } break;
        }
    }

    error_push((Error){
        .type = ERR_TYPE_RUNTIME,
        .value = ERR_PC_OOB,
        .line = frame->chunk->lines[0].line,
        .col = frame->chunk->lines[0].col,
    });
}

#undef FIXNUM_ARITHMETIC_OP
#undef FIXNUM_COMPARE_OP
#undef LOGIC_OP

#endif // BDL_VM_H