#include <lcq/pit/runtime/eval.h>

#include <stdio.h>

/* add a new stack frame to the vm (with a tag and callsite annotation) */
static void vm_push_code_func(pit_runtime *rt, pit_value code, pit_value tag, pit_annotation ann, pit_value bound) {
    pit_callstack_entry ent;
    ent.tag = tag;
    ent.ann = ann;
    ent.code = code;
    ent.bound = bound;
    if (pit_vec_push(pit_callstack_entry)(rt->callstack, ent) < 0) {
        pit_error(rt, "call stack overflow");
    }
}

/* add a new stack frame to the vm (with no annotation) */
static void vm_push_code(pit_runtime *rt, pit_value code) {
    pit_annotation ann;
    ann.line = -1;
    ann.column = -1;
    vm_push_code_func(rt, code, PIT_NIL, ann, PIT_NIL);
}

static void vm_push(pit_runtime *rt, pit_value v) {
    if (pit_vec_push(pit_value)(rt->result_stack, v) < 0) {
        pit_error(rt, "vm stack overflow");
    }
}

static pit_value vm_pop(pit_runtime *rt) {
    pit_value ret = PIT_NIL;
    if (pit_vec_pop(pit_value)(rt->result_stack, &ret) < 0) {
        pit_error(rt, "vm stack underflow");
    }
    return ret;
}

static void vm_call_lisp(pit_runtime *rt, pit_value tag, pit_value closure, pit_value args) {
    pit_value bound = PIT_NIL;
    pit_value env = pit_value_array_get(rt, closure, 0);
    pit_value anames = pit_value_array_get(rt, closure, 1);
    pit_value arg_rest_nm = pit_value_array_get(rt, closure, 2);
    pit_value body = pit_value_array_get(rt, closure, 3);
    if (rt->error != PIT_NIL) return;
    while (env != PIT_NIL) { /* first, bind all entries in the closure */
        pit_value b = pit_value_cons_car(rt, env);
        pit_value nm = pit_value_cons_car(rt, b);
        pit_symtab_bind(rt, nm, pit_value_cons_cdr(rt, b));
        bound = pit_value_cons(rt, nm, bound);
        env = pit_value_cons_cdr(rt, env);
    }
    while (anames != PIT_NIL) { /* bind all argument names to their values */
        pit_value nm = pit_value_cons_car(rt, anames);
        pit_value cell = pit_value_cell_new(rt, PIT_NIL);
        if (arg_rest_nm != PIT_NIL && pit_value_eq(nm, arg_rest_nm)) {
            pit_value_cell_set(rt, cell, args, nm);
            pit_symtab_bind(rt, nm, cell);
            break;
        } else {
            pit_value_cell_set(rt, cell, pit_value_cons_car(rt, args), nm);
            pit_symtab_bind(rt, nm, cell);
        }
        bound = pit_value_cons(rt, nm, bound);
        args = pit_value_cons_cdr(rt, args);
        anames = pit_value_cons_cdr(rt, anames);
    }
    pit_annotation ann;
    ann.line = -1;
    ann.column = -1;
    vm_push_code_func(rt, body, tag, ann, bound);
}

static bool vm_call(pit_runtime *rt, pit_value f, pit_value args) {
    char buf[256] = {0};
    if (pit_value_is_symbol(rt, f)) f = pit_symtab_fget(rt, f);
    switch (pit_value_sort(f)) {
    case PIT_VALUE_SORT_REF: {
        pit_value_heavy *h = pit_value_ref_deref(rt, pit_value_as_ref(rt, f));
        if (!h) { pit_error(rt, "bad ref for function"); return false; }
        switch (h->hsort) {
        case PIT_VALUE_HEAVY_SORT_NATIVEFUNC:
            vm_push(rt, h->in.nativefunc.f(rt, args, h->in.nativefunc.data));
            break;
        case PIT_VALUE_HEAVY_SORT_FUNC:
            vm_call_lisp(rt, h->in.func.nm, h->in.func.closure, args);
            break;
        default: {
            i64 end = pit_dump(rt, buf, sizeof(buf) - 1, f, true);
            buf[end] = 0;
            pit_error(rt, "attempted to apply non-function ref: %s", buf);
            return false;
        }
        }
        break;
    }
    default: {
        i64 end = pit_dump(rt, buf, sizeof(buf) - 1, f, true);
        buf[end] = 0;
        pit_error(rt, "attempted to apply non-function value: %s", buf);
        return false;
    }
    }
    return true;
}

static void vm_erase_empty(pit_runtime *rt) {
    pit_callstack_entry *ent = NULL;
    while ((ent = pit_vec_get(pit_callstack_entry)(rt->callstack, rt->callstack->next - 1))
        && ent->code == PIT_NIL) {
        pit_value bound = ent->bound;
        while (bound != PIT_NIL) { /* unbind everything we bound for the frame, in reverse */
            pit_symtab_unbind(rt, pit_value_cons_car(rt, bound));
            bound = pit_value_cons_cdr(rt, bound);
        }
        pit_vec_pop(pit_callstack_entry)(rt->callstack, NULL); /* pop the frame, and try again */
    }
}

/* run one instruction of the VM */
bool pit_vm_run_one(pit_runtime *rt) {
    if (rt->callstack->next < 1) return false;
    vm_erase_empty(rt);
    pit_callstack_entry *ent = pit_vec_get(pit_callstack_entry)(rt->callstack, rt->callstack->next - 1);
    if (ent == NULL) { pit_error(rt, "malformed call stack"); return false; }
    // for (i64 i = 0; i < rt->callstack->next; ++i) {
    //     pit_callstack_entry *e = pit_vec_get(pit_callstack_entry)(rt->callstack, i);
    //     fprintf(stderr, "frame: "); pit_dump_to_file(rt, stderr, e->code, false); fprintf(stderr, "\n");
    // }
    pit_value ins = pit_value_cons_car(rt, ent->code);
    // fprintf(stderr, "ins: "); pit_dump_to_file(rt, stderr, ins, false); fprintf(stderr, "\n");
    // fprintf(stderr, "foo: "); pit_dump_to_file(rt, stderr, pit_symtab_get_function_cell(rt, pit_symtab_intern_cstr(rt, "foo")), false); fprintf(stderr, "\n");
    if (ins == PIT_NIL) { pit_error(rt, "malformed vm instruction"); return false; }
    ent->code = pit_value_cons_cdr(rt, ent->code);
    pit_value op = pit_value_cons_car(rt, ins);
    if (pit_symtab_symbol_name_match_cstr(rt, op, "literal")) {
        /* push a lisp value to the vm stack */
        vm_push(rt, pit_value_cons_car(rt, pit_value_cons_cdr(rt, ins)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "lambda")) {
        /* pop an argument list, a free variable list, and a body block */
        /* create and push a function */
        /* notably, this builds the closure by recording the currently-bound cells for each free variable! */
        pit_value args = vm_pop(rt);
        pit_value freevars = vm_pop(rt);
        pit_value body = vm_pop(rt);
        vm_push(rt, pit_value_func_lambda(rt, args, freevars, body));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "drop")) {
        /* pop a value and ignore it */
        vm_pop(rt);
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "get")) {
        /* pop a symbol and look up its value */
        vm_push(rt, pit_symtab_get(rt, vm_pop(rt)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "fget")) {
        /* pop a symbol and look up its function value */
        vm_push(rt, pit_symtab_fget(rt, vm_pop(rt)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "if")) {
        /* pop a condition, a then-block, and an else-block */
        /* run the then-block if the condition is true, otherwise run the else-block */
        pit_value c = vm_pop(rt);
        pit_value t = vm_pop(rt);
        pit_value e = vm_pop(rt);
        vm_erase_empty(rt);
        vm_push_code(rt, c != PIT_NIL ? t : e);
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "apply")) {
        /* pop an arity n and a function, and then n arguments, and apply the function */
        i64 arity = pit_value_as_integer(rt, pit_value_cons_car(rt, pit_value_cons_cdr(rt, ins)));
        pit_value f = vm_pop(rt);
        pit_value args = PIT_NIL;
        while (arity-- > 0) args = pit_value_cons(rt, vm_pop(rt), args);
        vm_erase_empty(rt);
        if (!vm_call(rt, f, args)) return false;
    } else {
        pit_error(rt, "unknown vm operation");
        return false;
    }
    return true;
}

/* run the VM until evaluation of this code finishes, returning the result */
pit_value pit_vm_eval(pit_runtime *rt, pit_value v) {
    i64 start = rt->callstack->next;
    vm_push_code(rt, v);
    while (pit_vm_run_one(rt) && rt->callstack->next > start) {
        if (rt->expr_stack->next == 0 && rt->compilation_stack->next == 0) {
            pit_gc(rt); /* TODO hack to avoid running GC during macroexpansion */
        }
    }
    return vm_pop(rt);
}

/* run the VM until this function call finishes, returning the result */
pit_value pit_vm_apply(pit_runtime *rt, pit_value f, pit_value args) {
    i64 start = rt->callstack->next;
    vm_call(rt, f, args);
    while (pit_vm_run_one(rt) && rt->callstack->next > start);
    return vm_pop(rt);
}