#include "gc.h"

#include "function.h"
#include "hashtable.h"
#include "lisp_math.h"
#include "lisp_string.h"
#include "list.h"
#include "stack.h"
#include "symbol.h"

#include <gmp.h>
#include <stdlib.h>

bool lisp_doing_gc;
struct timespec total_gc_time;
size_t lisp_gc_count;

struct GCObjectList {
    LispVal *obj;
    struct GCObjectList *prev;
    struct GCObjectList *next;
};

#define FREE_OBJECTS_LIST_LIMIT 1024
static size_t free_objects_list_count;
static struct GCObjectList *free_objects_list;

static struct GCObjectList *black_objects;
static struct GCObjectList *gray_objects;
static struct GCObjectList *white_objects;
static struct GCObjectList *static_objects;

ObjectGCSet GC_BLACK = 0;
ObjectGCSet GC_GRAY = 1;
ObjectGCSet GC_WHITE = 2;

enum IncrementalGCSetp {
    GC_STEP_STATICS,
    GC_STEP_STACK,
    GC_STEP_HEAP,
    GC_STEP_FREE,
};

struct IncrementalGCState {
    enum IncrementalGCSetp step;
    struct GCObjectList *next_static;
};

static struct IncrementalGCState incremental_state = {
    .step = GC_STEP_STATICS,
    .next_static = NULL,
};

static ALWAYS_INLINE struct GCObjectList **HEAD_FOR_SET(ObjectGCSet set) {
    if (set == GC_BLACK) {
        return &black_objects;
    } else if (set == GC_GRAY) {
        return &gray_objects;
    } else if (set == GC_WHITE) {
        return &white_objects;
    } else {
        abort();
    }
}

static struct GCObjectList *alloc_gc_objects_list_node(void) {
    if (free_objects_list) {
        struct GCObjectList *to_return = free_objects_list;
        free_objects_list = free_objects_list->next;
        --free_objects_list_count;
        return to_return;
    } else {
        return lisp_malloc(sizeof(struct GCObjectList));
    }
}

static void unuse_gc_objects_list_node(struct GCObjectList *node) {
    node->next = free_objects_list;
    free_objects_list = node;
    ++free_objects_list_count;
}

void lisp_gc_register_object(void *val) {
    if (!OBJECTP(val)) {
        return;
    }
    LispObject *obj = val;
    obj->gc.is_static = false;
    obj->gc.set = GC_WHITE;
    struct GCObjectList *node = alloc_gc_objects_list_node();
    obj->gc.gc_node = node;
    node->prev = NULL;
    node->next = white_objects;
    if (node->next) {
        node->next->prev = node;
    }
    node->obj = val;
    white_objects = node;
}

void lisp_gc_register_static_object(void *val) {
    if (!OBJECTP(val)) {
        return;
    }
    LispObject *obj = val;
    obj->gc.is_static = true;
    struct GCObjectList *node = alloc_gc_objects_list_node();
    node->prev = NULL;
    node->next = static_objects;
    if (node->next) {
        node->next->prev = node;
    }
    node->obj = obj;
    static_objects = node;
    // reset incremental GC to ensure we scan the new static
    incremental_state.step = GC_STEP_STATICS;
    incremental_state.next_static = static_objects;
}

static void unregister_object_node(LispObject *obj) {
    struct GCObjectList *node = obj->gc.gc_node;
    if (!node->prev) {
        *HEAD_FOR_SET(obj->gc.set) = node->next;
    } else {
        node->prev->next = node->next;
    }
    if (node->next) {
        node->next->prev = node->prev;
    }
}

void gc_move_to_set(void *val, ObjectGCSet new_set) {
    if (!OBJECTP(val)) {
        return;
    }
    LispObject *obj = val;
    if (obj->gc.set != new_set) {
        struct GCObjectList *node = obj->gc.gc_node;
        unregister_object_node(obj);
        obj->gc.set = new_set;
        node->prev = NULL;
        node->next = *HEAD_FOR_SET(new_set);
        if (node->next) {
            node->next->prev = node;
        }
        *HEAD_FOR_SET(new_set) = node;
    }
}

void gc_mark_stack_for_rescan(void) {
    if (incremental_state.step > GC_STEP_STACK) {
        incremental_state.step = GC_STEP_STACK;
    }
}

static void free_object(LispVal *val) {
    // This is called on non-white objects during cleanup! Don't assert
    // OBJECT_GC_SET_P!
    assert(!OBJECT_LOWEST_LOCAL_REFERENCE(val));
    switch (((LispObject *) val)->type) {
    case TYPE_HASH_TABLE: {
        LispHashTable *ht = val;
        lisp_free(ht->data);
        break;
    }
    case TYPE_STRING: {
        LispString *str = val;
        if (str->owned) {
            lisp_free(str->data);
        }
        break;
    }
    case TYPE_VECTOR: {
        LispVector *vec = val;
        lisp_free(vec->data);
        break;
    }
    case TYPE_GMP:
        mpz_clear(((LispGmp *) val)->val);
        break;
    case TYPE_CONS:
    case TYPE_SYMBOL:
    case TYPE_FUNCTION:
        // nothing to do
        break;
    case TYPE_FIXNUM:
    case TYPE_FLOAT:
    default:
        abort();
    }
    unregister_object_node(val);
    unuse_gc_objects_list_node(((LispObject *) val)->gc.gc_node);
    lisp_release_object(val);
}

static inline void make_gray_if_white(LispVal *val) {
    if (OBJECTP(val) && OBJECT_GC_SET_P(val, GC_WHITE)) {
        gc_move_to_set(val, GC_GRAY);
    }
}

static void mark_object(LispVal *val) {
    // check for null for newly constructed objects
    if (!val || !OBJECTP(val) || OBJECT_GC_SET_P(val, GC_BLACK)) {
        return;
    }
    switch (((LispObject *) val)->type) {
    case TYPE_CONS:
        make_gray_if_white(((LispCons *) val)->car);
        make_gray_if_white(((LispCons *) val)->cdr);
        break;
    case TYPE_SYMBOL: {
        LispSymbol *sym = val;
        make_gray_if_white(sym->name);
        make_gray_if_white(SYMBOL_VALUE(sym));
        make_gray_if_white(sym->function);
        make_gray_if_white(sym->plist);
        break;
    }
    case TYPE_VECTOR: {
        LispVector *vec = val;
        for (size_t i = 0; i < vec->length; ++i) {
            make_gray_if_white(vec->data[i]);
        }
        break;
    }
    case TYPE_HASH_TABLE: {
        HT_FOREACH_INDEX(val, i) {
            make_gray_if_white(HASH_KEY(val, i));
            make_gray_if_white(HASH_VALUE(val, i));
        }
        break;
    }
    case TYPE_FUNCTION: {
        LispFunction *fobj = val;
        make_gray_if_white(fobj->name);
        make_gray_if_white(fobj->docstr);
        make_gray_if_white(fobj->args.req);
        make_gray_if_white(fobj->args.opt);
        make_gray_if_white(fobj->args.kw);
        make_gray_if_white(fobj->args.rest);
        break;
    }
    case TYPE_STRING:
    case TYPE_GMP:
        // no held refs
        break;
    case TYPE_FIXNUM:
    case TYPE_FLOAT:
    default:
        abort();
    }
    gc_move_to_set(val, GC_BLACK);
}

static inline size_t saturating_dec(size_t *restrict limit, size_t amount) {
    if (amount >= *limit) {
        *limit = 0;
    } else {
        *limit -= amount;
    }
    return *limit;
}

static void mark_statics(size_t *restrict limit) {
    struct GCObjectList *node = incremental_state.next_static;
    while (node && saturating_dec(limit, 1)) {
        mark_object(node->obj);
        node = node->next;
    }
    // we processed the whole list, move to the next step
    if (!node) {
        incremental_state.next_static = static_objects;
        incremental_state.step = GC_STEP_STACK;
    }
}

// This mark_stack_local_refs and mark_stack_frame mark the whole frame,
// ignoring limit. However, they update limit with how many objects the marked.
static void mark_stack_local_refs(struct LocalReferences *restrict refs,
                                  size_t *restrict limit) {
    size_t full_blocks = refs->num_refs / LOCAL_REFERENCES_BLOCK_LENGTH;
    size_t last_block_len = refs->num_refs % LOCAL_REFERENCES_BLOCK_LENGTH;
    for (size_t i = 0; i < full_blocks; ++i) {
        for (size_t j = 0; j < LOCAL_REFERENCES_BLOCK_LENGTH; ++j) {
            mark_object(refs->blocks[i]->refs[j]);
        }
    }
    for (size_t i = 0; i < last_block_len; ++i) {
        mark_object(refs->blocks[full_blocks]->refs[i]);
    }
    saturating_dec(limit, refs->num_refs);
}

static void mark_stack_frame(StackFrame *frame, size_t *restrict limit) {
    switch (frame->kind) {
    case STACK_FRAME_LOCAL_REFERENCES:
        mark_stack_local_refs(&frame->local_references, limit);
        break;
    case STACK_FRAME_CALL:
        mark_object(frame->call.name);
        mark_object(frame->call.args);
        mark_object(frame->call.fobj);
        saturating_dec(limit, 3);
        break;
    case STACK_FRAME_UNWIND_PROTECT:
        // nothing to do
        break;
    case STACK_FRAME_HANDLER_BIND:
        mark_object(frame->handler_bind.exceptions);
        saturating_dec(limit, 1);
        break;
    case STACK_FRAME_DYNAMIC_BINDING:
        mark_object(frame->dynamic_binding.symbol);
        mark_object(frame->dynamic_binding.old_value);
        saturating_dec(limit, 2);
        break;
    case STACK_FRAME_BLOCK:
        mark_object(frame->block.tag);
        // leave protecting value up to internal-return-from
        saturating_dec(limit, 1);
        break;
    }
}

static void mark_the_stack(size_t *restrict limit) {
    size_t i;
    for (i = 0; i < the_stack.depth && *limit; ++i) {
        if (!the_stack.frames[i].marked) {
            mark_stack_frame(&the_stack.frames[i], limit);
            the_stack.frames[i].marked = true;
        }
    }
    if (i == the_stack.depth) {
        incremental_state.step = GC_STEP_HEAP;
    }
}

static void unmark_the_stack(void) {
    for (size_t i = 0; i < the_stack.depth; ++i) {
        the_stack.frames[i].marked = false;
    }
}

static void mark_gray_objects(size_t *restrict limit) {
    while (gray_objects && saturating_dec(limit, 1)) {
        mark_object(gray_objects->obj);
    }
    if (!gray_objects) {
        incremental_state.step = GC_STEP_FREE;
    }
}

static void swap_white_black_sets(void) {
    struct GCObjectList *tmp_node = white_objects;
    white_objects = black_objects;
    black_objects = tmp_node;
    ObjectGCSet tmp_id = GC_WHITE;
    GC_WHITE = GC_BLACK;
    GC_BLACK = tmp_id;
}

static void maybe_free_some_object_list_nodes(void) {
    while (free_objects_list_count > FREE_OBJECTS_LIST_LIMIT) {
        struct GCObjectList *to_free = free_objects_list;
        free_objects_list = free_objects_list->next;
        lisp_free(to_free);
        --free_objects_list_count;
    }
}

static void gc_sweep_objects(size_t *restrict limit) {
    while (white_objects && saturating_dec(limit, 1)) {
        assert(OBJECT_GC_SET_P(white_objects->obj, GC_WHITE));
        free_object(white_objects->obj);
    }
    // reset the gc
    if (!white_objects) {
        swap_white_black_sets();
        maybe_free_some_object_list_nodes();
        unmark_the_stack();
        incremental_state.step = GC_STEP_STATICS;
    }
}

void lisp_gc_yield(struct timespec *restrict time_took, bool full) {
    lisp_doing_gc = true;
    struct timespec start_time;
    clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
    size_t limit = full ? SIZE_MAX : LISP_GC_INCREMENTAL_COUNT;
    while (limit) {
        // there are more GRAY objects, mark them before we sweep
        if (incremental_state.step == GC_STEP_FREE && gray_objects) {
            incremental_state.step = GC_STEP_HEAP;
        }
        switch (incremental_state.step) {
        case GC_STEP_STATICS:
            mark_statics(&limit);
            break;
        case GC_STEP_STACK:
            mark_the_stack(&limit);
            break;
        case GC_STEP_HEAP:
            mark_gray_objects(&limit);
            break;
        case GC_STEP_FREE:
            gc_sweep_objects(&limit);
            // force being done
            limit = 0;
            break;
        }
    }
    struct timespec end_time;
    clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &end_time);
    struct timespec backup_time_took;
    if (!time_took) {
        time_took = &backup_time_took;
    }
    sub_timespecs(&end_time, &start_time, time_took);
    add_timespecs(time_took, &total_gc_time, &total_gc_time);
    lisp_doing_gc = false;
    ++lisp_gc_count;
}

void lisp_gc_teardown(void) {
    while (white_objects) {
        free_object(white_objects->obj);
    }
    while (gray_objects) {
        free_object(gray_objects->obj);
    }
    while (black_objects) {
        free_object(black_objects->obj);
    }
    while (static_objects) {
        struct GCObjectList *next = static_objects->next;
        free(static_objects);
        static_objects = next;
    }
    while (free_objects_list) {
        struct GCObjectList *next = free_objects_list->next;
        free(free_objects_list);
        free_objects_list = next;
    }
}