#include "error.h"
#include "op.h"
#include "unit.h"
#include "vmstate.h"
#include "vmval.h"

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

static inline int64_t sximm(uint32_t in) {
    if (in & 0x800000) {
        return in | 0xFFFFFFFFFF000000;
    } else {
        return in;
    }
}

static int vm_push_context(struct vm_state *vm) {
    int res;
    if ((res = stack_push(&vm->call_stack, vm->lp)) != 0) {
        return res;
    }
    if ((res = stack_push(&vm->call_stack, vm->fp)) != 0) {
        return res;
    }
    if ((res = stack_push(&vm->call_stack, vm->ip)) != 0) {
        return res;
    }
    return 0;
}

static int vm_pop_context(struct vm_state *vm) {
    int res;
    if ((res = stack_pop(&vm->call_stack, &vm->ip)) != 0) {
        return res;
    }
    if ((res = stack_pop(&vm->call_stack, &vm->fp)) != 0) {
        return res;
    }
    if ((res = stack_pop(&vm->call_stack, &vm->lp)) != 0) {
        return res;
    }
    return 0;
}

int vm_state_init(struct vm_state *vm, size_t stack_size) {
    int res;

    if ((res = stack_init(&vm->data_stack, stack_size)) != 0) {
        return res;
    }
    if ((res = stack_init(&vm->call_stack, 1024)) != 0) {
        return res;
    }

    vector_init(&vm->units, sizeof(struct vm_unit));

    vm->lp = 0;
    vm->fp = 0;
    vm->ip = 0;

    return 0;
}

void vm_state_free(struct vm_state *vm) {
    struct vm_unit *unit;
    stack_free(&vm->data_stack);
    stack_free(&vm->call_stack);
    for (size_t i = 0; i < vm->units.size; ++i) {
        unit = vector_ref(&vm->units, i);
        unit_free(unit);
    }
    vector_free(&vm->units);
}

struct vm_unit *vm_add_unit(struct vm_state *vm, size_t global_pool_size) {
    struct vm_unit *unit = vector_append(&vm->units);
    if (!unit) {
        return unit;
    }

    unit->is_loaded = 0;
    unit->global_pool = calloc(sizeof(uint64_t), global_pool_size);
    unit->global_pool_size = global_pool_size;

    vector_init(&unit->ref_table, sizeof(struct vm_ref_entry));
    vector_init(&unit->functions, sizeof(struct vm_func_entry));
    return unit;
}

// Call an entry from function table of a unit
int vm_call_unit_index(struct vm_state *vm, size_t lib_index, size_t fn_index) {
    struct vm_func_entry *func;
    struct vm_unit *unit;
    int res;

    unit = vector_ref(&vm->units, lib_index);
    func = vector_ref(&unit->functions, fn_index);

    if (func->argc > MAXARG) {
        return -ERR_RANGE;
    }

    if ((res = vm_push_context(vm)) != 0) {
        return res;
    }

    vm->lp = lib_index;
    vm->fp = fn_index;
    vm->ip = 0;

    for (size_t i = 0; i < func->argc; ++i) {
        if ((res = stack_pop(&vm->data_stack, &func->arg_regs[i])) != 0) {
            return res;
        }
    }

    return 0;
}

int vm_call_ref(struct vm_state *vm, struct vm_value *ref) {
    if (ref->type != VT_FUNC) {
        return -ERR_OPERAND_TYPE;
    }
    return vm_call_unit_index(vm, ref->v_func.lib_index, ref->v_func.fn_index);
}

static int vm_read_ext_ref(struct vm_state *vm, struct vm_ref_entry *ref, uint64_t *w) {
    struct vm_unit *unit;
    int res;
    unit = vector_ref(&vm->units, ref->unit_index);
    if (!unit->is_loaded) {
        // Before trying to access target value, eval the unit first
        if ((res = vm_eval_unit(vm, ref->unit_index)) != 0) {
            return res;
        }
    }
    assert(ref->ref_index < unit->global_pool_size);
    if (ref->ref_index >= unit->global_pool_size) {
        return -ERR_RANGE;
    }
    *w = unit->global_pool[ref->ref_index];
    return 0;
}

static int vm_eval_debug(struct vm_state *vm, struct vm_func_entry *func, uint32_t arg) {
    (void) func;
    uint64_t w0;
    int res;
    switch (arg) {
    case 0x01:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        printf("trace: ");
        vm_print(w0);
        printf("\n");
        vm_unref(w0);
        return 0;
    default:
        return -ERR_OPCODE_UNDEFINED;
    }
}

int vm_eval_step(struct vm_state *vm) {
    uint64_t w0, w1;
    int64_t sw0, sw1;
    ssize_t ii0;
    size_t i0;
    struct vm_ref_entry *r0;
    struct vm_value *v0;
    int res;

    assert(vm->lp < vm->units.size);
    struct vm_unit *unit = vector_ref(&vm->units, vm->lp);
    assert(vm->fp < unit->functions.size);
    struct vm_func_entry *func = vector_ref(&unit->functions, vm->fp);
    uint32_t opcode = func->bytecode[vm->ip++];

    printf("%02zx:%02zx:%02zx: %08x\n", vm->lp, vm->fp, vm->ip, opcode);
    switch (opcode >> 24) {
    case OP_ADD:
        if ((res = vm_pop_integer(vm, &sw0)) != 0) {
            return res;
        }
        if ((res = vm_pop_integer(vm, &sw1)) != 0) {
            return res;
        }
        return vm_push_integer(vm, sw0 + sw1);
    case OP_SUB:
        if ((res = vm_pop_integer(vm, &sw0)) != 0) {
            return res;
        }
        if ((res = vm_pop_integer(vm, &sw1)) != 0) {
            return res;
        }
        return vm_push_integer(vm, sw0 - sw1);
    case OP_MUL:
        if ((res = vm_pop_integer(vm, &sw0)) != 0) {
            return res;
        }
        if ((res = vm_pop_integer(vm, &sw1)) != 0) {
            return res;
        }
        return vm_push_integer(vm, sw0 * sw1);
    case OP_NOT:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        res = vm_push_bool(vm, null_q(w0));
        vm_unref(w0);
        return res;
    case OP_EQ:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        if ((res = stack_pop(&vm->data_stack, &w1)) != 0) {
            return res;
        }
        if (ref_q(w0) || ref_q(w1)) {
            // TODO what do?
            return -ERR_OPERAND_TYPE;
        }
        res = vm_push_bool(vm, w0 == w1);
        vm_unref(w0);
        vm_unref(w1);
        return res;
//
    case OP_LDNIL:
        return stack_push(&vm->data_stack, FLAG_REF);
    case OP_LDI:
        return vm_push_integer(vm, sximm(opcode & 0xFFFFFF));
    case OP_CAR:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        if (!pair_q(w0)) {
            return -ERR_OPERAND_TYPE;
        }
        vm_ref(getref(w0)->v_cons.fat_ar);
        res = stack_push(&vm->data_stack, getref(w0)->v_cons.fat_ar);
        vm_unref(w0);
        return res;
    case OP_CDR:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        if (!pair_q(w0)) {
            return -ERR_OPERAND_TYPE;
        }
        vm_ref(getref(w0)->v_cons.fat_dr);
        res = stack_push(&vm->data_stack, getref(w0)->v_cons.fat_dr);
        vm_unref(w0);
        return res;
    case OP_CONS:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        if ((res = stack_pop(&vm->data_stack, &w1)) != 0) {
            return res;
        }
        v0 = vm_cons(w0, w1);
        vm_value_ref(v0);
        res = vm_push_ref(vm, v0);
        vm_unref(w0);
        vm_unref(w1);
        return res;
//
    case OP_LDARG:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= MAXARG) {
            return -ERR_RANGE;
        }
        w0 = func->arg_regs[i0];
        vm_ref(w0);
        return stack_push(&vm->data_stack, w0);
    case OP_STARG:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= MAXARG) {
            return -ERR_RANGE;
        }
        vm_unref(func->arg_regs[i0]);
        return stack_pop(&vm->data_stack, &func->arg_regs[i0]);
    case OP_LDG:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= unit->global_pool_size) {
            return -ERR_RANGE;
        }
        w0 = unit->global_pool[i0];
        vm_ref(w0);
        return stack_push(&vm->data_stack, w0);
    case OP_STG:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= unit->global_pool_size) {
            return -ERR_RANGE;
        }
        vm_unref(unit->global_pool[i0]);
        return stack_pop(&vm->data_stack, &unit->global_pool[i0]);
    case OP_LDF:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= unit->functions.size) {
            return -ERR_RANGE;
        }
        v0 = vm_func(vm->lp, i0);
        vm_value_ref(v0);
        return vm_push_ref(vm, v0);
    case OP_STL:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= func->local_count) {
            return -ERR_RANGE;
        }
        vm_unref(func->local_regs[i0]);
        return stack_pop(&vm->data_stack, &func->local_regs[i0]);
    case OP_LDL:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= func->local_count) {
            return -ERR_RANGE;
        }
        w0 = func->local_regs[i0];
        vm_ref(w0);
        return stack_push(&vm->data_stack, w0);
//
    case OP_ISZ:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        res = vm_push_bool(vm, null_q(w0));
        vm_unref(w0);
        return res;
//
    case OP_XCALL:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= unit->ref_table.size) {
            return -ERR_RANGE;
        }
        r0 = vector_ref(&unit->ref_table, i0);
        if (r0->flags & REF_NATIVE) {
            int (*func) (struct vm_state *) = (void *) r0->ref_native;
            if (!func) {
                return -ERR_OPERAND_TYPE;
            }
            return func(vm);
        } else {
            if ((res = vm_read_ext_ref(vm, r0, &w0)) != 0) {
                return res;
            }
            if (!func_q(w0)) {
                return -ERR_OPERAND_TYPE;
            }
            return vm_call_ref(vm, getref(w0));
        }
    case OP_GCALL:
        i0 = opcode & 0xFFFFFF;
        if (i0 >= unit->global_pool_size) {
            return -ERR_RANGE;
        }
        w0 = unit->global_pool[i0];
        if (!func_q(w0)) {
            return -ERR_OPERAND_TYPE;
        }
        return vm_call_ref(vm, getref(w0));
    case OP_CALL:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        if (!func_q(w0)) {
            return -ERR_OPERAND_TYPE;
        }
        res = vm_call_ref(vm, getref(w0));
        vm_unref(w0);
        return res;
    case OP_BF:
        if ((res = stack_pop(&vm->data_stack, &w0)) != 0) {
            return res;
        }
        ii0 = sximm(opcode & 0xFFFFFF);
        if (null_q(w0)) {
            if (ii0 == 0) {
                return -ERR_RANGE;
            }
            vm->ip += ii0 - 1;
        }
        vm_unref(w0);
        return 0;
    case OP_JMP:
        ii0 = sximm(opcode & 0xFFFFFF);
        if (ii0 == 0) {
            return -ERR_RANGE;
        }
        vm->ip += ii0 - 1;
        return 0;
    case OP_RET:
        // Drop references to arguments
        for (size_t i = 0; i < func->argc; ++i) {
            vm_unref(func->arg_regs[i]);
        }
        return vm_pop_context(vm);
//
    case OP_DEBUG:
        return vm_eval_debug(vm, func, opcode & 0xFFFFFF);
    default:
        fprintf(stderr, "Undefined opcode: %02x\n", opcode >> 24);
        return -ERR_OPCODE_UNDEFINED;
    }
}

int vm_eval_unit(struct vm_state *vm, size_t index) {
    struct vm_unit *unit;
    size_t csp;
    int res;

    unit = vector_ref(&vm->units, index);

    csp = vm->call_stack.sp;
    vm_call_unit_index(vm, index, 0);

    // Run until return
    while (vm->call_stack.sp != csp) {
        if ((res = vm_eval_step(vm)) != 0) {
            return res;
        }
    }
    unit->is_loaded = 1;

    return 0;
}