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

#include <stdio.h>

static void 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;
    }
    }
}

static pit_value free_vars(pit_runtime *rt, pit_value initial_bound, pit_value body) {
    i64 expr_stack_reset = rt->expr_stack->next;
    pit_value ret = PIT_NIL;
    if (pit_vec_push(pit_value)(rt->expr_stack, pit_value_cons(rt, initial_bound, body)) < 0) {
        pit_error(rt, "free variable search stack overflow");
        return PIT_NIL;
    }
    while (rt->expr_stack->next > expr_stack_reset) {
        pit_value boundscur, bound, cur;
        if (pit_vec_pop(pit_value)(rt->expr_stack, &boundscur) < 0) {
            pit_error(rt, "free variable search stack underflow");
            return PIT_NIL;
        }
        bound = pit_value_cons_car(rt, boundscur);
        cur = pit_value_cons_cdr(rt, boundscur);
        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_value fargs = pit_value_cons_cdr(rt, cur);
            if (is_symbol && pit_symtab_symbol_name_match_cstr(rt, fsym, "lambda")) {
                pit_value new_bound = pit_value_list_append(rt, pit_value_cons_car(rt, fargs), bound);
                fargs = pit_value_cons_cdr(rt, fargs);
                while (fargs != PIT_NIL) {
                    if (pit_vec_push(pit_value)(rt->expr_stack, pit_value_cons(rt, new_bound, pit_value_cons_car(rt, fargs))) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                    fargs = pit_value_cons_cdr(rt, fargs);
                }
            } else if (is_symbol && pit_symtab_symbol_name_match_cstr(rt, fsym, "quote")) {
                /* don't look inside quote!
                   if we add other special forms, make sure to consider them here if necessary! */
            } else {
                while (fargs != PIT_NIL) {
                    if (pit_vec_push(pit_value)(rt->expr_stack, pit_value_cons(rt, bound, pit_value_cons_car(rt, fargs))) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                    fargs = pit_value_cons_cdr(rt, fargs);
                }
                if (!is_symbol) {
                    if (pit_vec_push(pit_value)(rt->expr_stack, pit_value_cons(rt, bound, fsym)) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                }
            }
        } else if (pit_value_is_symbol(rt, cur)) {
            if (pit_value_list_contains_eq(rt, cur, bound) == PIT_NIL) {
                ret = pit_value_cons(rt, cur, ret);
            }
        }
    }
    rt->expr_stack->next = expr_stack_reset;
    return ret;
}

static pit_value lambda(pit_runtime *rt, pit_value args, pit_value body) {
    pit_value expanded = pit_macroexpand(rt, pit_value_cons(rt, pit_symtab_intern_cstr(rt, "progn"), body));
    fprintf(stderr, "lambda: "); pit_dump_to_file(rt, stderr, expanded, false); fprintf(stderr, "\n");
    return pit_value_list(rt, 4,
        pit_symtab_intern_cstr(rt, "lambda"),
        args,
        free_vars(rt, args, expanded),
        pit_compile(rt, expanded)
    );
}

static void c_now(pit_runtime *rt, pit_value v) {
    pit_traversal_push_value(rt, rt->traversal, v);
}

static void c_eval(pit_runtime *rt, pit_value e) {
    if (pit_vec_push(pit_value)(rt->expr_stack, e) < 0)
        pit_error(rt, "evaluation stack overflow");
}

pit_value pit_compile(pit_runtime *rt, pit_value top) {
    char buf[256] = {0};
    pit_value ret = PIT_NIL;
    i64 expr_stack_reset = rt->expr_stack->next;
    i64 traversal_reset = rt->traversal->next;
    fprintf(stderr, "compile: "); pit_dump_to_file(rt, stderr, top, false); fprintf(stderr, "\n");
    c_eval(rt, top);
    /* 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);
                call_special_form(rt, f, args);
            } else if (is_symbol && pit_symtab_is_symbol_macro(rt, fsym)) { /* macros */
                i64 end = pit_dump(rt, buf, sizeof(buf) - 1, fsym, true);
                buf[end] = 0;
                pit_error(rt, "encountered an unexpanded macro while compiling: %s", buf);
            } else { /* normal functions */
                pit_value args = pit_value_cons_cdr(rt, cur);
                i64 argcount = 0;
                while (args != PIT_NIL) {
                    // fprintf(stderr, "push1: "); pit_dump_to_file(rt, stderr, pit_value_cons_car(rt, args), false); fprintf(stderr, "\n");
                    c_eval(rt, pit_value_cons_car(rt, args));
                    args = pit_value_cons_cdr(rt, args);
                    argcount += 1;
                }
                if (!is_symbol) {
                    // fprintf(stderr, "push2: "); pit_dump_to_file(rt, stderr, fsym, false); fprintf(stderr, "\n");
                    c_eval(rt, fsym);
                }
                c_now(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "apply"), pit_value_integer_new(rt, argcount)));
                if (is_symbol) {
                    c_now(rt, pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "fget")));
                    c_now(rt, 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) {
                c_now(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), cur));
            } else {
                c_now(rt, pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "get")));
                c_now(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), cur));
            }
        } else { /* other expressions evaluate to themselves! */
            c_now(rt, 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;
    }
}

static pit_value impl_sf_quote(pit_runtime *rt, pit_value args, void *data) {
    (void) data;
    pit_traversal_push_value(rt, rt->traversal,
        pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), pit_value_cons_car(rt, args))
    );
    return PIT_NIL;
}
static pit_value impl_sf_if(pit_runtime *rt, pit_value args, void *data) {
    (void) data;
    pit_value c = pit_value_cons_car(rt, args);
    args = pit_value_cons_cdr(rt, args);
    pit_value t = pit_value_cons_car(rt, args);
    args = pit_value_cons_cdr(rt, args);
    pit_value e = pit_value_cons_car(rt, args);
    c_now(rt, pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "if")));
    c_now(rt, lambda(rt, PIT_NIL, pit_value_list(rt, 1, t)));
    c_now(rt, lambda(rt, PIT_NIL, pit_value_list(rt, 1, e)));
    c_eval(rt, c);
    return PIT_NIL;
}
static pit_value impl_sf_lambda(pit_runtime *rt, pit_value args, void *data) {
    (void) data;
    pit_value as = pit_value_cons_car(rt, args);
    pit_value body = pit_value_cons_cdr(rt, args);
    c_now(rt, lambda(rt, as, body));
    return PIT_NIL;
}

void pit_compile_install_special_forms(pit_runtime *rt) {
    pit_symtab_sfset(rt, pit_symtab_intern_cstr(rt, "quote"), pit_value_nativefunc_new(rt, impl_sf_quote));
    pit_symtab_sfset(rt, pit_symtab_intern_cstr(rt, "if"), pit_value_nativefunc_new(rt, impl_sf_if));
    pit_symtab_sfset(rt, pit_symtab_intern_cstr(rt, "lambda"), pit_value_nativefunc_new(rt, impl_sf_lambda));
}