#include "stack.h"

#include "function.h"
#include "hashtable.h"
#include "lisp_string.h"
#include "list.h"
#include "memory.h"
#include "print.h"
#include "symbol.h"

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

struct LispStack the_stack;

void lisp_init_stack(void) {
    the_stack.max_depth = LISP_STACK_MAX_DEPTH;
    the_stack.depth = 0;
    the_stack.frames = lisp_malloc(sizeof(StackFrame) * the_stack.max_depth);
}

static inline void
maybe_clear_object_local_reference(StackFrame *restrict frame, LispVal *obj) {
    if (OBJECT_LOWEST_LOCAL_REFERENCE(obj) == frame) {
        SET_OBJECT_LOWEST_LOCAL_REFERENCE(obj, NULL);
    }
}

static void teardown_local_references(StackFrame *restrict frame) {
    assert(frame->kind == STACK_FRAME_LOCAL_REFERENCES);
    struct LocalReferences *restrict refs = &frame->local_references;
    size_t last_block_size = refs->num_refs % LOCAL_REFERENCES_BLOCK_LENGTH;
    size_t num_full_blocks = refs->num_refs / LOCAL_REFERENCES_BLOCK_LENGTH;
    for (size_t i = 0; i < num_full_blocks; ++i) {
        for (size_t j = 0; j < LOCAL_REFERENCES_BLOCK_LENGTH; ++j) {
            assert(OBJECTP(refs->blocks[i]->refs[j]));
            maybe_clear_object_local_reference(frame, refs->blocks[i]->refs[j]);
        }
        lisp_free(refs->blocks[i]);
    }
    for (size_t i = 0; i < last_block_size; ++i) {
        assert(OBJECTP(refs->blocks[num_full_blocks]->refs[i]));
        maybe_clear_object_local_reference(
            frame, refs->blocks[num_full_blocks]->refs[i]);
    }
    lisp_free(refs->blocks[num_full_blocks]);
    lisp_free(refs->blocks);
}

void lisp_teardown_stack(void) {
    for (size_t i = 0; i < the_stack.depth; ++i) {
        if (the_stack.frames[i].kind == STACK_FRAME_LOCAL_REFERENCES) {
            teardown_local_references(&the_stack.frames[i]);
        }
    }
    lisp_free(the_stack.frames);
}

static ALWAYS_INLINE StackFrame *PUSH_NEW_FRAME(enum StackFrameKind kind) {
    if (the_stack.depth == LISP_STACK_SOFT_MAX_DEPTH) {
        lisp_signal(Qexcessive_lisp_nesting_error, Qnil);
    }
    StackFrame *last_refs =
        the_stack.depth ? the_stack.frames[the_stack.depth - 1].last_references
                        : NULL;
    StackFrame *frame = &the_stack.frames[the_stack.depth++];
    frame->kind = kind;
    frame->marked = false;
    frame->last_references = last_refs;
    gc_mark_stack_for_rescan();
    return frame;
}

void push_call_frame(LispVal *name, LispVal *args) {
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_CALL);
    frame->call.name = name;
    frame->call.args = args;
    frame->call.evaled_args = false;
}

void set_stack_evaluated_args(StackFrame *restrict frame, LispVal *fobj,
                              LispVal *args) {
    gc_mark_stack_for_rescan();
    assert(frame->kind == STACK_FRAME_CALL);
    frame->call.fobj = fobj;
    frame->call.args = args;
    frame->call.evaled_args = true;
}

void push_unwind_protect_frame(void (*handler)(void *user_ptr), void *user_ptr,
                               void (*cleanup_user_ptr)(void *data)) {
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_UNWIND_PROTECT);
    frame->unwind_protect.handler = handler;
    frame->unwind_protect.user_ptr = user_ptr;
    frame->unwind_protect.cleanup_user_ptr = cleanup_user_ptr;
}

void push_handler_bind_frame(LispVal *exceptions,
                             void (*handler)(LispVal *name, LispVal *data,
                                             void *user_ptr),
                             void *user_ptr, void (*cleanup_user_ptr)(void *)) {
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_HANDLER_BIND);
    frame->handler_bind.enabled = true;
    frame->handler_bind.exceptions = exceptions;
    frame->handler_bind.handler = handler;
    frame->handler_bind.user_ptr = user_ptr;
    frame->handler_bind.cleanup_user_ptr = cleanup_user_ptr;
}

void push_local_reference_frame(void) {
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_LOCAL_REFERENCES);
    struct LocalReferences *refs = &frame->local_references;
    refs->num_refs = 0;
    refs->num_blocks = 1;
    refs->blocks = lisp_malloc(sizeof(struct LocalReferencesBlock *));
    refs->blocks[0] = lisp_malloc(sizeof(struct LocalReferencesBlock));
    frame->last_references = frame;
}

static void store_local_reference_in_frame(StackFrame *restrict frame,
                                           LispVal *obj) {
    struct LocalReferences *restrict refs = &frame->local_references;
    size_t num_full_blocks = refs->num_refs / LOCAL_REFERENCES_BLOCK_LENGTH;
    if (num_full_blocks == refs->num_blocks) {
        refs->blocks =
            lisp_realloc(refs->blocks, sizeof(struct LocalReferencesBlock *)
                                           * ++refs->num_blocks);
        refs->blocks[refs->num_blocks - 1] =
            lisp_malloc(sizeof(struct LocalReferencesBlock));
        refs->blocks[refs->num_blocks - 1]->refs[0] = obj;
        refs->num_refs += 1;
    } else {
        refs->blocks[num_full_blocks]
            ->refs[refs->num_refs++ % LOCAL_REFERENCES_BLOCK_LENGTH] = obj;
    }
    SET_OBJECT_LOWEST_LOCAL_REFERENCE(obj, frame);
    // mark the frame for rescan
    frame->marked = false;
    gc_mark_stack_for_rescan();
}

void add_local_reference_no_recurse(StackFrame *restrict frame, LispVal *obj) {
    if (!OBJECTP(obj)) {
        return;
    }
    frame = frame->last_references;
    assert(frame->kind == STACK_FRAME_LOCAL_REFERENCES);
    StackFrame *current_frame = OBJECT_LOWEST_LOCAL_REFERENCE(obj);
    if (!current_frame || current_frame > frame) {
        store_local_reference_in_frame(frame, obj);
    }
}

static LispVal *next_local_reference(StackFrame *restrict frame,
                                     size_t *restrict i) {
    if (*i >= frame->local_references.num_refs) {
        return NULL;
    }
    size_t block_idx = *i / LOCAL_REFERENCES_BLOCK_LENGTH;
    size_t small_idx = *i % LOCAL_REFERENCES_BLOCK_LENGTH;
    LispVal *obj = frame->local_references.blocks[block_idx]->refs[small_idx];
    ++*i;
    return obj;
}

static inline void
add_local_ref_if_not_seen_no_recurse(StackFrame *restrict frame,
                                     LispVal *seen_objs, LispVal *obj) {
    if (NILP(Fgethash(seen_objs, obj, Qnil))) {
        add_local_reference_no_recurse(frame, obj);
        Fputhash(seen_objs, obj, Qt);
    }
}

static inline void
add_local_refs_for_object_sub_vals(StackFrame *restrict frame,
                                   LispVal *seen_objs, LispVal *val) {
    switch (((LispObject *) val)->type) {
    case TYPE_CONS:
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                             ((LispCons *) val)->car);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                             ((LispCons *) val)->cdr);
        break;
    case TYPE_SYMBOL: {
        LispSymbol *sym = val;
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, sym->name);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                             SYMBOL_VALUE(sym));
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, sym->function);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, sym->plist);
        break;
    }
    case TYPE_VECTOR: {
        LispVector *vec = val;
        for (size_t i = 0; i < vec->length; ++i) {
            add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                                 vec->data[i]);
        }
        break;
    }
    case TYPE_HASH_TABLE: {
        HT_FOREACH_INDEX(val, i) {
            add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                                 HASH_KEY(val, i));
            add_local_ref_if_not_seen_no_recurse(frame, seen_objs,
                                                 HASH_VALUE(val, i));
        }
        break;
    }
    case TYPE_FUNCTION: {
        LispFunction *fobj = val;
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->name);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->docstr);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->args.req);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->args.opt);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->args.kw);
        add_local_ref_if_not_seen_no_recurse(frame, seen_objs, fobj->args.rest);
        break;
    }
    case TYPE_STRING:
    case TYPE_GMP:
        // no held refs
        break;
    case TYPE_FIXNUM:
    case TYPE_FLOAT:
    default:
        abort();
    }
}

void add_local_reference(StackFrame *restrict frame, LispVal *obj) {
    if (!OBJECTP(obj)) {
        return;
    }
    frame = frame->last_references;
    assert(frame->kind == STACK_FRAME_LOCAL_REFERENCES);
    add_local_reference_no_recurse(frame, obj);
    LispVal *seen_objs = make_hash_table_no_gc(Qnil, Qnil);
    Fputhash(seen_objs, obj, Qt);
    size_t i = frame->local_references.num_refs - 1;
    LispVal *cur;
    while ((cur = next_local_reference(frame, &i))) {
        add_local_refs_for_object_sub_vals(frame, seen_objs, cur);
    }
    release_hash_table_no_gc(seen_objs);
}

void push_dynamic_binding(LispVal *name, LispVal *new_value) {
    assert(SYMBOLP(name));
    LispVal *old_val = SYMBOL_VALUE(name);
    SET_SYMBOL_VALUE(name, new_value); // will throw if name is const
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_DYNAMIC_BINDING);
    frame->dynamic_binding.symbol = name;
    frame->dynamic_binding.old_value = old_val;
}

void set_lexical_variable(LispVal *name, LispVal *value) {
    assert(SYMBOLP(name));
    if (CONST_VALUE_P(name)) {
        signal_value_constant(name);
    }
    if (DYNAMIC_SYMBOL_P(name)) {
        SET_SYMBOL_VALUE(name, value);
    } else {
        Vlexical_environment = Fplist_put(Vlexical_environment, name, value);
    }
}

void new_lexical_variable(LispVal *name, LispVal *value) {
    assert(SYMBOLP(name));
    if (CONST_VALUE_P(name)) {
        signal_value_constant(name);
    }
    if (DYNAMIC_SYMBOL_P(name)) {
        push_dynamic_binding(name, value);
    } else {
        Vlexical_environment = CONS(name, CONS(value, Vlexical_environment));
    }
}

void push_block_frame(LispVal *tag, jmp_buf *target,
                      LispVal *volatile *value_ptr, StackFrame *unwind_to) {
    StackFrame *frame = PUSH_NEW_FRAME(STACK_FRAME_BLOCK);
    frame->block.tag = tag;
    frame->block.target = target;
    frame->block.value_ptr = value_ptr;
    frame->block.unwind_to = unwind_to;
}

void unwind_to(StackFrame *frame) {
    while (the_stack.depth && &the_stack.frames[the_stack.depth - 1] > frame) {
        StackFrame *restrict top = &the_stack.frames[--the_stack.depth];
        switch (top->kind) {
        case STACK_FRAME_DYNAMIC_BINDING:
            if (top->dynamic_binding.symbol == Qlexical_environment) {
                Vlexical_environment = top->dynamic_binding.old_value;
            } else {
                SET_SYMBOL_VALUE(top->dynamic_binding.symbol,
                                 top->dynamic_binding.old_value);
            }
            break;
        case STACK_FRAME_UNWIND_PROTECT:
            top->unwind_protect.handler(top->unwind_protect.user_ptr);
            if (top->unwind_protect.cleanup_user_ptr) {
                top->unwind_protect.cleanup_user_ptr(
                    top->unwind_protect.user_ptr);
            }
            break;
        case STACK_FRAME_LOCAL_REFERENCES:
            teardown_local_references(top);
            break;
        case STACK_FRAME_HANDLER_BIND:
        case STACK_FRAME_CALL:
        case STACK_FRAME_BLOCK:
            // nothing to do
            break;
        }
    }
}

static bool frame_handles_exception(StackFrame *restrict frame,
                                    LispVal *exception_name) {
    assert(frame->kind == STACK_FRAME_HANDLER_BIND);
    DOLIST(fe, frame->handler_bind.exceptions) {
        if (!NILP(Fcondition_subclass_p(exception_name, fe))) {
            return true;
        }
    }
    return false;
}

static StackFrame *find_exception_handler(StackFrame *from,
                                          LispVal *exception_name) {
    for (ptrdiff_t i = from - the_stack.frames; i >= 0; --i) {
        StackFrame *restrict top = &the_stack.frames[i];
        if (top->kind == STACK_FRAME_HANDLER_BIND && top->handler_bind.enabled
            && frame_handles_exception(top, exception_name)) {
            return top;
        }
    }
    return NULL;
}

static void do_run_handler_cases(LispVal *name, LispVal *data,
                                 StackFrame *start) {
    while ((start = find_exception_handler(start, name))) {
        start->handler_bind.enabled = false;
        start->handler_bind.handler(name, data, start->handler_bind.user_ptr);
        start->handler_bind.enabled = true;
        if (start->handler_bind.cleanup_user_ptr) {
            start->handler_bind.cleanup_user_ptr(start->handler_bind.user_ptr);
        }
    }
}

static noreturn void top_of_stack_exception_handler(LispVal *name,
                                                    LispVal *data) {
    fprintf(stderr,
            "An exception has propagated to the top of the lisp stack!\n");
    fprintf(stderr, "Name: ");
    debug_print(stderr, name);
    fprintf(stderr, "\nData: ");
    debug_print(stderr, data);
    fprintf(stderr, "\nLisp will now exit...\n");
    abort();
}

noreturn void lisp_signal(LispVal *name, LispVal *data) {
    if (NILP(Fcondition_class_p(name))) {
        signal_type_error(name, LIST(Qcondition_class));
    }
    CHECK_LISTP(data);
    do_run_handler_cases(name, data, LISP_STACK_REF());
    unwind_to(NULL);
    top_of_stack_exception_handler(name, data);
}

DEFUN(backtrace, "backtrace", (void), "()", "") {
    LispVal *out = Qnil;
    for (size_t i = 0; i < the_stack.depth; ++i) {
        StackFrame *restrict frame = &the_stack.frames[i];
        if (frame->kind == STACK_FRAME_CALL) {
            LispVal *name = frame->call.name;
            // fobj is NULL (not Qnil) if the arguments haven't been evaluated
            LispVal *fobj = frame->call.fobj ? frame->call.fobj : Qnil;
            LispVal *did_eval_args = frame->call.evaled_args ? Qt : Qnil;
            LispVal *args = frame->call.args;
            out = CONS(LIST(name, fobj, did_eval_args, args), out);
        }
    }
    return out;
}

DEFINE_SYMBOL(no_such_block_error, "no-such-block-error");
DEFINE_CONDITION_CLASS(no_such_block_error, error);
DEFINE_SYMBOL(block_out_of_scope_error, "block-out-of-scope-error");
DEFINE_CONDITION_CLASS(block_out_of_scope_error, error);
DEFINE_SYMBOL(excessive_lisp_nesting_error, "excessive-lisp-nesting-error");
DEFINE_CONDITION_CLASS(excessive_lisp_nesting_error, error);