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

/* helper functions for manipulating the compiler state */
static void c_now(pit_runtime *rt, pit_value v) {
    /* push a literal value to the traversal */
    pit_traversal_push_value(rt, rt->traversal, v);
}
static void c_compile(pit_runtime *rt, pit_value e) {
    /* schedule an expression to be compiled */
    /* the resulting entries will be pushed to the traversal subsequently */
    pit_compilation_push_value(rt, rt->compilation_stack, e);
}
static void c_drop(pit_runtime *rt) {
    /* schedule a drop to be compiled */
    pit_compilation_push_drop(rt, rt->compilation_stack);
}
static void c_code(pit_runtime *rt, pit_value e) {
    /* schedule an expression to be compiled, and then captured as code */
    /* this will result in a literal code value being included in the program, */
    /* rather than the code itself (as in c_compile) */
    pit_compilation_push_end_code(rt, rt->compilation_stack);
    pit_compilation_push_value(rt, rt->compilation_stack, e);
    pit_compilation_push_begin_code(rt, rt->compilation_stack);
}

/* given an expression and some initially bound variables (e.g. arguments), */
/* compute the free variables in the expression */
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)) {
            /* if the expression is a list, determine if it is a normal application */
            pit_value fsym = pit_value_cons_car(rt, cur);
            bool is_symbol = pit_value_is_symbol(rt, fsym);
            bool is_special_form = is_symbol && pit_symtab_is_symbol_special_form(rt, fsym);
            pit_value fargs = pit_value_cons_cdr(rt, cur);
            /* first, consider special forms that have special behavior with regard to binding */
            if (is_special_form && pit_symtab_symbol_name_match_cstr(rt, fsym, "quote")) {
                /* don't look inside quote!
                   NOTICE if we add other special forms, make sure to consider them here if necessary! */
            } else if (is_special_form && pit_symtab_symbol_name_match_cstr(rt, fsym, "lambda")) {
                /* the variables from the lambda argument list are now bound */
                pit_value new_bound = pit_value_list_append(rt, pit_value_cons_car(rt, fargs), bound);
                fargs = pit_value_cons_cdr(rt, fargs);
                /* consider all body forms within the lambda with these new bindings in mind */ 
                while (fargs != PIT_NIL) {
                    pit_value bindings_expr_pair = pit_value_cons(rt, new_bound, pit_value_cons_car(rt, fargs));
                    if (pit_vec_push(pit_value)(rt->expr_stack, bindings_expr_pair) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                    fargs = pit_value_cons_cdr(rt, fargs);
                }
            } else {
                /* otherwise, this is a normal application form */
                /* first consider all arguments to the application */
                while (fargs != PIT_NIL) {
                    pit_value bindings_expr_pair = pit_value_cons(rt, bound, pit_value_cons_car(rt, fargs));
                    if (pit_vec_push(pit_value)(rt->expr_stack, bindings_expr_pair) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                    fargs = pit_value_cons_cdr(rt, fargs);
                }
                if (!is_symbol) {
                    /* ... and then consider the function being applied, if it is not a symbol */
                    pit_value bindings_expr_pair = pit_value_cons(rt, bound, fsym);
                    if (pit_vec_push(pit_value)(rt->expr_stack, bindings_expr_pair) < 0) {
                        pit_error(rt, "free variable search stack overflow");
                        return PIT_NIL;
                    }
                }
            }
        } else if (pit_value_is_symbol(rt, cur)) {
            /* if the expression is a symbol, check if it's free! */
            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 void compile_special_form(pit_runtime *rt, pit_value f, pit_value args) {
    char buf[256] = {0};
    if (pit_symtab_symbol_name_match_cstr(rt, f, "quote")) {
        /* quote compiles to a literal instruction */
        c_now(rt,
            pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), pit_value_cons_car(rt, args))
        );
    } else if (pit_symtab_symbol_name_match_cstr(rt, f, "if")) {
        /* if compiles the condition, and then code for the branches, and then the if instruction */
        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_code(rt, e);
        c_code(rt, t);
        c_compile(rt, c);
    } else if (pit_symtab_symbol_name_match_cstr(rt, f, "progn")) {
        /* progn simply compiles each of its argument forms in order */
        /* after every form but the final form, it inserts a drop to ignore the return value */
        while (args != PIT_NIL) {
            pit_value form = pit_value_cons_car(rt, args);
            args = pit_value_cons_cdr(rt, args);
            c_compile(rt, form);
            if (args != PIT_NIL) {
                c_drop(rt);
            }
        }
    } else if (pit_symtab_symbol_name_match_cstr(rt, f, "lambda")) {
        pit_value arglist = pit_value_cons_car(rt, args);
        pit_value body = pit_value_cons(rt, pit_symtab_intern_cstr(rt, "progn"), pit_value_cons_cdr(rt, args));
        c_now(rt, pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "lambda")));
        c_now(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), arglist));
        c_now(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), free_vars(rt, arglist, body)));
        c_code(rt, body);
    } else {
        i64 end = pit_dump(rt, buf, sizeof(buf) - 1, f, true);
        buf[end] = 0;
        pit_error(rt, "unknown special form: %s", buf);
    }
}

static void compile_output_instruction(pit_runtime *rt, pit_value ins) {
    /* write an instruction to the code block at the top of the compilation stack */
    pit_compilation_entry b;
    if (pit_vec_pop(pit_compilation_entry)(rt->compilation_stack, &b) < 0) {
        pit_error(rt, "compilation block stack underflow");
        return;
    }
    pit_value new = pit_value_cons(rt, ins, b.in.value);
    pit_compilation_push_value(rt, rt->compilation_stack, new);
}

static void compile_interpret_traversal_entry(pit_runtime *rt, pit_traversal_entry *ent) {
    if (ent == NULL) {
        pit_error(rt, "evaluation traversal invalid");
        return;
    }
    switch (ent->sort) {
    case PIT_TRAVERSAL_ENTRY_VALUE: {
        /* normal traversal entries just cons values to the code block at the top of the compilation stack */ 
        compile_output_instruction(rt, ent->in.value);
        return;
    }
    case PIT_TRAVERSAL_ENTRY_BEGIN_CODE: {
        /* starting a new code block pushes a new list to the top of the compilation stack */
        pit_compilation_push_value(rt, rt->compilation_stack, PIT_NIL);
        return;
    }
    case PIT_TRAVERSAL_ENTRY_END_CODE: {
        /* ending a code block pops a code block, and pushes it as a literal in the frame below */
        pit_compilation_entry b;
        if (pit_vec_pop(pit_compilation_entry)(rt->compilation_stack, &b) < 0) {
            pit_error(rt, "compilation block stack underflow");
            return;
        }
        compile_output_instruction(rt, pit_value_list(rt, 2, pit_symtab_intern_cstr(rt, "literal"), b.in.value));
        return;
    }
    default:
        pit_error(rt, "unknown traversal entry");
        return;
    } 
}

pit_value pit_compile(pit_runtime *rt, pit_value top) {
    /* compilation proceeds in a somewhat unintuitive manner */
    /* typically, we would simply recursively traverse the tree, accumulating the output code */
    /* (this also has the advantage of allowing the results of the recursive compilation call to */
    /* be used in nonstandard contexts, such as compiling the body of lambdas. */
    /* however, in pit we would like to avoid this sort of recursion, as we'd like to be able to */
    /* make claims about the interpreter's C stack usage (and also, in general, it seems quite */
    /* nice to avoid C stack usage growing with expression size!) */
    /* as such, we must explicitly represent the control/data flow that would otherwise be handled */
    /* implicitly by the recursion. we do this by breaking the compilation step into several steps. */
    /* first, we traverse the expression using a stack (rt->compilation_stack). during this step, we */
    /* push subexpression to be compiled to the stack, and append "generated code" to an intermediate */
    /* structure (rt->traversal). this intermediate structure contains "operations" to later perform */
    /* - think of it as the "continuation" of the code that produced it. finally, these intermediate */
    /* operations are "interpreted" to generate the real compiled code, which is a Lisp list of instructions */
    /* each instruction is itself a list of opcode and operands (see eval.c for more information on this) */
    char buf[256] = {0};
    i64 compilation_stack_reset = rt->compilation_stack->next;
    i64 traversal_reset = rt->traversal->next;
    /* before we do anything else, we fully macroexpand the input */
    top = pit_macroexpand(rt, top);
    c_compile(rt, top);
    while (rt->compilation_stack->next > compilation_stack_reset) {
        pit_compilation_entry curent;
        if (rt->error != PIT_NIL) goto err;
        if (pit_vec_pop(pit_compilation_entry)(rt->compilation_stack, &curent) < 0) {
            pit_error(rt, "compilation stack underflow");
            goto err;
        }
        switch (curent.sort) {
        case PIT_COMPILATION_ENTRY_VALUE: {
            pit_value cur = curent.in.value;
            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);
                if (is_symbol && pit_symtab_is_symbol_special_form(rt, fsym)) { /* special forms */
                    pit_value args = pit_value_cons_cdr(rt, cur);
                    compile_special_form(rt, fsym, 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) {
                        c_compile(rt, pit_value_cons_car(rt, args));
                        args = pit_value_cons_cdr(rt, args);
                        argcount += 1;
                    }
                    if (!is_symbol) {
                        c_compile(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));
            }
            break;
        }
        case PIT_COMPILATION_ENTRY_DROP: {
            /* a BEGIN_CODE is a marker indicating that a new "block" of code has started */
            c_now(rt, pit_value_list(rt, 1, pit_symtab_intern_cstr(rt, "drop")));
            break;
        }
        case PIT_COMPILATION_ENTRY_BEGIN_CODE: {
            /* a BEGIN_CODE is a marker indicating that a new "block" of code has started */
            pit_traversal_push_begin_code(rt, rt->traversal);
            break;
        }
        case PIT_COMPILATION_ENTRY_END_CODE: {
            /* an END_CODE is a marker indicating that the current "block" of code has finished */
            pit_traversal_push_end_code(rt, rt->traversal);
            break;
        }
        }
    }
    /* now that we've built up the traversal, we iterate over it to build up the final program */
    /* the "current" code is at the top of the compilation stack (which we now repurpose to this end) */
    /* (note that this usage of the compilation state is entirely separate from the previous usage!) */
    rt->compilation_stack->next = compilation_stack_reset;
    pit_compilation_push_value(rt, rt->compilation_stack, PIT_NIL);
    for (i64 idx = traversal_reset; idx < rt->traversal->next; idx++) {
        pit_traversal_entry *ent = pit_vec_get(pit_traversal_entry)(rt->traversal, idx);
        if (rt->error != PIT_NIL) goto end;
        compile_interpret_traversal_entry(rt, ent);
    }
    goto end;
err: {
        /* in case of error, we should always consider the stack empty, so we don't return something spurious */
        rt->compilation_stack->next = compilation_stack_reset;
    }
end: {
        /* finally, we are done (possibly due to an error) */
        pit_value ret = PIT_NIL;
        if (rt->compilation_stack->next > compilation_stack_reset) {
            /* if the compilation stack is not empty, the top is our compiled code! */
            pit_compilation_entry ent;
            pit_vec_pop(pit_compilation_entry)(rt->compilation_stack, &ent);
            ret = ent.in.value;
        }
        if (rt->compilation_stack->next != compilation_stack_reset) {
            /* if there was more than one entry on the compilation stack, this is a bug */
            pit_error(rt, "compilation did not reduce stack to single code! this is a bug!");
        }
        rt->compilation_stack->next = compilation_stack_reset;
        rt->traversal->next = traversal_reset;
        return ret;
    }
}

/* mark the symbols for the special forms as such */
void pit_compile_install_special_forms(pit_runtime *rt) {
    pit_symtab_symbol_mark_special_form(rt, pit_symtab_intern_cstr(rt, "quote"));
    pit_symtab_symbol_mark_special_form(rt, pit_symtab_intern_cstr(rt, "if"));
    pit_symtab_symbol_mark_special_form(rt, pit_symtab_intern_cstr(rt, "progn"));
    pit_symtab_symbol_mark_special_form(rt, pit_symtab_intern_cstr(rt, "lambda"));
}