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

#include <stdio.h>

void pit_vm_call_special_form(pit_runtime *rt, pit_value f, pit_value args) {
    char buf[256] = {0};
    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 special form"); return; }
        switch (h->hsort) {
        case PIT_VALUE_HEAVY_SORT_NATIVEFUNC:
            h->in.nativefunc.f(rt, args, h->in.nativefunc.data);
            break;
        default: {
            i64 end = pit_dump(rt, buf, sizeof(buf) - 1, f, true);
            buf[end] = 0;
            pit_error(rt, "attempted to apply non-nativefunc special form: %s", buf);
            return;
        }
        }
        break;
    }
    default: {
        i64 end = pit_dump(rt, buf, sizeof(buf) - 1, f, true);
        buf[end] = 0;
        pit_error(rt, "attempted to apply non-function special form: %s", buf);
        return;
    }
    }
}

pit_value pit_vm_compile(pit_runtime *rt, pit_value top) {
    pit_value ret = PIT_NIL;
    i64 expr_stack_reset = rt->expr_stack->next;
    i64 traversal_reset = rt->traversal->next;
    if (pit_vec_push(pit_value)(rt->expr_stack, top) < 0)
        pit_error(rt, "evaluation stack overflow");
    /* first, convert the expression tree into "polish notation" in traversal */
    while (rt->expr_stack->next > expr_stack_reset) {
        pit_value cur = PIT_NIL;
        if (rt->error != PIT_NIL) goto end;
        if (pit_vec_pop(pit_value)(rt->expr_stack, &cur) < 0)
            pit_error(rt, "evaluation stack underflow");
        fprintf(stderr, "cur: "); pit_dump_to_file(rt, stderr, cur, false); fprintf(stderr, "\n");
        if (pit_value_is_cons(rt, cur)) {
            pit_value fsym = pit_value_cons_car(rt, cur);
            bool is_symbol = pit_value_is_symbol(rt, fsym);
            pit_annotation *ann = pit_annotation_get(rt, pit_value_as_ref(rt, cur));
            if (is_symbol && pit_symtab_is_symbol_special_form(rt, fsym)) { /* special forms */
                pit_value f = pit_symtab_fget(rt, fsym);
                pit_value args = pit_value_cons_cdr(rt, cur);
                pit_vm_call_special_form(rt, f, args);
            } else if (is_symbol && pit_symtab_is_symbol_macro(rt, fsym)) { /* macros */
                pit_error(rt, "encountered a macro while evaluating");
            } else { /* normal functions */
                pit_value args = pit_value_cons_cdr(rt, cur);
                i64 argcount = 0;
                while (args != PIT_NIL) {
                    if (pit_vec_push(pit_value)(rt->expr_stack, pit_value_cons_car(rt, args)) < 0)
                        pit_error(rt, "evaluation stack overflow");
                    args = pit_value_cons_cdr(rt, args);
                    argcount += 1;
                }
                if (!is_symbol) {
                    if (pit_vec_push(pit_value)(rt->expr_stack, fsym) < 0)
                        pit_error(rt, "evaluation stack overflow");
                }
                pit_traversal_push_value(rt, rt->traversal,
                    pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "apply"), pit_value_integer_new(rt, argcount))
                );
                if (is_symbol) {
                    pit_traversal_push_value(rt, rt->traversal,
                        pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "fget"))
                    );
                    pit_traversal_push_value(rt, rt->traversal,
                        pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), fsym)
                    );
                }
            }
        } else if (pit_value_is_symbol(rt, cur)) { /* unquoted symbols: variable lookup */
            pit_symtab_entry *ent = pit_symtab_lookup(rt, cur);
            if (ent->is_keyword) {
                pit_traversal_push_value(rt, rt->traversal,
                    pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), cur)
                );
            } else {
                pit_traversal_push_value(rt, rt->traversal,
                    pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "get"))
                );
                pit_traversal_push_value(rt, rt->traversal,
                    pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), cur)
                );
            }
        } else { /* other expressions evaluate to themselves! */
            pit_traversal_push_value(rt, rt->traversal,
                pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), cur)
            );
        }
    }
    for (i64 idx = traversal_reset; idx < rt->traversal->next; idx++) {
        pit_traversal_entry *ent = pit_vec_get(pit_traversal_entry)(rt->traversal, idx);
        if (ent == NULL) pit_error(rt, "evaluation traversal invalid");
        if (rt->error != PIT_NIL) goto end;
        switch (ent->sort) {
        case PIT_TRAVERSAL_ENTRY_VALUE: {
            ret = pit_value_cons(rt, ent->in.value, ret);
            break;
        }
        default:
            pit_error(rt, "unknown traversal entry");
            ret = PIT_NIL;
            goto end;
        }
    }
end: {
        rt->expr_stack->next = expr_stack_reset;
        rt->traversal->next = traversal_reset;
        fprintf(stderr, "compiled: "); pit_dump_to_file(rt, stderr, ret, false); fprintf(stderr, "\n");
        return ret;
    }
}

/* add a new stack frame to the vm (with a tag and callsite annotation) */
void pit_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;
    fprintf(stderr, "push code: "); pit_dump_to_file(rt, stderr, code, false); fprintf(stderr, "\n");
    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) */
void pit_vm_push_code(pit_runtime *rt, pit_value code) {
    pit_annotation ann;
    ann.line = -1;
    ann.column = -1;
    pit_vm_push_code_func(rt, code, PIT_NIL, ann, PIT_NIL);
}

void pit_vm_push(pit_runtime *rt, pit_value v) {
    fprintf(stderr, "push: "); pit_dump_to_file(rt, stderr, v, false); fprintf(stderr, "\n");
    if (pit_vec_push(pit_value)(rt->result_stack, v) < 0) {
        pit_error(rt, "vm stack overflow");
    }
}

pit_value pit_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");
    }
    fprintf(stderr, "pop: "); pit_dump_to_file(rt, stderr, ret, false); fprintf(stderr, "\n");
    return ret;
}

void pit_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;
    pit_vm_push_code_func(rt, body, tag, ann, bound);
}

void pit_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; }
        switch (h->hsort) {
        case PIT_VALUE_HEAVY_SORT_NATIVEFUNC:
            pit_vm_push(rt, h->in.nativefunc.f(rt, args, h->in.nativefunc.data));
            break;
        case PIT_VALUE_HEAVY_SORT_FUNC:
            pit_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;
        }
        }
        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;
    }
    }
}

/* run one instruction of the VM */
bool pit_vm_run_one(pit_runtime *rt) {
    char buf[256] = {0};
    if (rt->callstack->next < 1) return false;
    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;
    }
    pit_value ins = pit_value_cons_car(rt, ent->code);
    if (ins == PIT_NIL) {
        pit_error(rt, "malformed vm instruction");
        return false;
    }
    pit_value bound = ent->bound;
    pit_value rest = pit_value_cons_cdr(rt, ent->code);
    bool final = false;
    if (rest == PIT_NIL) {
        final = true;
        pit_vec_pop(pit_callstack_entry)(rt->callstack, NULL);
    } else ent->code = rest;
    fprintf(stderr, "ins: "); pit_dump_to_file(rt, stderr, ins, false); fprintf(stderr, "\n");
    pit_value op = pit_value_cons_car(rt, ins);
    if (pit_symtab_symbol_name_match_cstr(rt, op, "literal")) {
        pit_vm_push(rt, pit_value_cons_car(rt, pit_value_cons_cdr(rt, ins)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "get")) {
        pit_vm_push(rt, pit_symtab_get(rt, pit_vm_pop(rt)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "fget")) {
        pit_vm_push(rt, pit_symtab_fget(rt, pit_vm_pop(rt)));
    } else if (pit_symtab_symbol_name_match_cstr(rt, op, "apply")) {
        i64 arity = pit_value_as_integer(rt, pit_value_cons_car(rt, pit_value_cons_cdr(rt, ins)));
        pit_value f = pit_vm_pop(rt);
        pit_value args = PIT_NIL;
        while (arity-- > 0) args = pit_value_cons(rt, pit_vm_pop(rt), args);
        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 PIT_NIL; }
            switch (h->hsort) {
            case PIT_VALUE_HEAVY_SORT_NATIVEFUNC:
                pit_vm_push(rt, h->in.nativefunc.f(rt, args, h->in.nativefunc.data));
                break;
            case PIT_VALUE_HEAVY_SORT_FUNC:
                pit_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;
        }
        }
    } else {
        pit_error(rt, "unknown vm operation");
        return false;
    }
    if (final) { /* if we have finished this stack frame */
        fprintf(stderr, "unbinding: "); pit_dump_to_file(rt, stderr, bound, false); fprintf(stderr, "\n");
        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);
        }
    }
    return true;
}

/* run the VM until evaluation finishes, returning the result */
pit_value pit_vm_eval(pit_runtime *rt, pit_value v) {
    pit_vm_push_code(rt, v);
    while (pit_vm_run_one(rt));
    pit_vm_pop(rt);
}

pit_value pit_vm_apply(pit_runtime *rt, pit_value f, pit_value args) {
    pit_vm_call(rt, f, args);
    while (pit_vm_run_one(rt));
    pit_vm_pop(rt);
}