#include "gc.h"

#include "function.h"
#include "hashtable.h"
#include "lisp.h"
#include "list.h"
#include "stack.h"

#include <stdlib.h>

bool lisp_doing_gc;
struct timespec total_gc_time;
size_t total_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 *grey_objects;
static struct GCObjectList *white_objects;
static struct GCObjectList *static_objects;

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

static ALWAYS_INLINE struct GCObjectList **HEAD_FOR_SET(ObjectGCSet set) {
    if (set == GC_BLACK) {
        return &black_objects;
    } else if (set == GC_GREY) {
        return &grey_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;
}

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

static void free_object(LispVal *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_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_grey_if_while(LispVal *val) {
    if (OBJECTP(val) && OBJECT_GC_SET_P(val, GC_WHITE)) {
        gc_move_to_set(val, GC_GREY);
    }
}

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

static void mark_statics(void) {
    for (struct GCObjectList *node = static_objects; node; node = node->next) {
        mark_object(node->obj);
    }
}

static void mark_stack_local_refs(struct LocalReferences *restrict refs) {
    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]);
    }
}

static void mark_stack_frame(struct StackFrame *frame) {
    mark_object(frame->name);
    mark_object(frame->args);
    mark_object(frame->fobj);
    mark_object(frame->lexenv);
    mark_stack_local_refs(&frame->local_refs);
}

static void mark_and_compact_the_stack(void) {
    mark_object(the_stack.nogc_retval);
    size_t i;
    for (i = 0; i < the_stack.depth; ++i) {
        mark_stack_frame(&the_stack.frames[i]);
    }
    for (; i < the_stack.first_clear_local_refs; ++i) {
        compact_stack_frame(&the_stack.frames[i]);
    }
    the_stack.first_clear_local_refs = the_stack.depth;
}

static void mark_grey_objects(void) {
    while (grey_objects) {
        mark_object(grey_objects->obj);
    }
}

static void gc_sweep_objects(void) {
    while (white_objects) {
        free_object(white_objects->obj);
    }
}

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

void lisp_gc_now(struct timespec *restrict time_took) {
    lisp_doing_gc = true;
    struct timespec start_time;
    clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &start_time);
    mark_statics();
    mark_object(obarray);
    mark_and_compact_the_stack();
    mark_grey_objects();
    gc_sweep_objects();
    maybe_free_some_object_list_nodes();
    swap_white_black_sets();
    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);
    ++total_gc_count;
    lisp_doing_gc = false;
}