#include "gc.h"

Environment *
alloc_env(void) {
    if (array_size(gc.free_envs.offsets) == 0) {
        mark_and_sweep();
        if (array_size(gc.free_envs.offsets) == 0) {
            fprintf(stderr, "NO MORE ENV MEMORY AVAILABLE WHERE IS YOUR GOD NOW MWAHAHA\n");
            dump_gc();
            exit(EXIT_FAILURE);
            // TODO: grow heap tables.
        }
    }
    size_t slot = gc.free_envs.offsets[gc.free_envs.position++];
    array_head(gc.free_envs.offsets)->size--;
    return &gc.envs[slot];
}

void
push_root(Object *obj) {
    array_push(gc.roots, obj);
}

Object *
pop_root(void) {
    return array_pop(gc.roots);
}

void
push_active_env(Environment *env) {
    array_push(gc.active_envs, env);
}

Environment *
pop_active_env(void) {
    return array_pop(gc.active_envs);
}

void
gc_init(void) {
    gc = (GC){0};

    array_init(gc.objects, GC_OBJS_CAP);
    array_init(gc.roots, GC_ROOTS_CAP);
    array_init(gc.active_envs, GC_ACTIVE_ENVS_CAP);
    array_init(gc.envs, GC_ENVS_CAP);
    array_init(gc.free_objects.offsets, GC_OBJS_CAP);
    array_init(gc.free_envs.offsets, GC_ENVS_CAP);

    // The free list stores the offset from the initial position for all
    // available slots.
    for (size_t i = 0; i < GC_OBJS_CAP; i++) {
        array_push(gc.free_objects.offsets, i);
    }
    for (size_t i = 0; i < GC_ENVS_CAP; i++) {
        array_push(gc.free_envs.offsets, i);
    }
}

void
mark_environment(Environment *env) {
    if (env == NULL || env->marked) {
        return;
    }
    env->marked = true;
    HashTablePair *pairs = env->table->pairs;
    for (size_t i = 0; i < array_cap(pairs); i++) {
        if (pairs[i].key != NULL) {
            mark_obj(pairs[i].key);
            mark_obj(pairs[i].value);
        }
    }
}

void
mark_obj(Object *obj) {
    if (obj->marked) {
        return;
    }
    obj->marked = true;
    if (obj->type == OBJ_TYPE_PAIR) {
        mark_obj(obj->car);
        mark_obj(obj->cdr);
    }
    if (obj->type == OBJ_TYPE_LAMBDA) {
        mark_obj(obj->params);
        mark_obj(obj->body);
        mark_environment(obj->env);
    }
}

void
mark_and_sweep(void) {
    // Mark.
    for (size_t i = 0; i < array_size(gc.active_envs); i++) {
        mark_environment(gc.active_envs[i]);
    }
    for (size_t i = 0; i < array_size(gc.roots); i++) {
        mark_obj(gc.roots[i]);
    }

    // Reset the free list.
    gc.free_objects.position = 0;
    array_head(gc.free_objects.offsets)->size = 0;
    gc.free_envs.position = 0;
    array_head(gc.free_envs.offsets)->size = 0;

    // Sweep.
    for (size_t i = 0; i < array_cap(gc.objects); i++) {
        Object *obj = &gc.objects[i];
        if (!obj->marked) {
            // Free heap allocated memory for this object if needed.
            if (obj->type == OBJ_TYPE_SYMBOL) {
                array_free(obj->symbol);
            } else if (obj->type == OBJ_TYPE_STRING) {
                array_free(obj->string);
            }
            gc.free_objects.offsets[array_head(gc.free_objects.offsets)->size++] = i;
        }
        obj->marked = false;
    }
    for (size_t i = 0; i < array_cap(gc.envs); i++) {
        Environment *env = &gc.envs[i];
        if (!env->marked) {
            ht_free(env->table);
            gc.free_envs.offsets[array_head(gc.free_envs.offsets)->size++] = i;
        }
        env->marked = false;
    }
}

void
dump_gc(void) {
        printf("-------------- ROOTS -------------- \n");
        for (size_t i = 0; i < array_size(gc.roots); i++) {
            display(gc.roots[i]);
            printf("\n");
        }
        printf("--------- OBJECTS (TOP 20) -------- \n");
        for (size_t i = 0; i < 20; i++) {
            printf("i: %ld -> ", i);
            Object *obj = &gc.objects[i];
            display(obj);
            bool is_free = false;
            for (size_t j = 0; j < array_cap(gc.objects); j++) {
                if (gc.free_objects.offsets[j] == i) {
                    is_free = true;
                    break;
                }
            }
            if (is_free) {
                printf(" [FREE]");
            }
            printf("\n");
        }
        printf("-------------- MISC --------------- \n");
        printf("gc.roots.size: %ld\n", array_size(gc.roots));
        printf("gc.roots.cap: %ld\n", array_size(gc.roots));
        printf("gc.active_envs.size: %ld\n", array_size(gc.active_envs));
        printf("gc.active_envs.cap: %ld\n", array_cap(gc.active_envs));
        printf("gc.obj_cap: %ld\n", array_cap(gc.objects));
        printf("gc.free_objects.size: %ld\n", array_size(gc.free_objects.offsets));
        printf("gc.free_objects.cap: %ld\n", array_cap(gc.free_objects.offsets));
        printf("gc.free_objects.position: %ld\n", gc.free_objects.position);
        printf("array_size(gc.free_envs.offsets): %ld\n", array_size(gc.free_envs.offsets));
        printf("gc.free_envs.cap: %ld\n", array_cap(gc.free_envs.offsets));
        printf("gc.free_envs.position: %ld\n", gc.free_envs.position);
        printf("gc.envs.size: %ld\n", array_size(gc.envs));
        printf("gc.envs.cap: %ld\n", array_cap(gc.envs));
}

Object *
alloc_object(ObjectType type) {
    if (array_head(gc.free_objects.offsets)->size == 0) {
        mark_and_sweep();
        if (array_head(gc.free_objects.offsets)->size == 0) {
            fprintf(stderr, "NO MORE OBJ MEMORY AVAILABLE WHERE IS YOUR GOD NOW MWAHAHA\n");
            dump_gc();
            exit(EXIT_FAILURE);
            // TODO: grow heap tables.
            // NOTE: When growing the tables, we WILL lose the pointer
            // references! Should we work with offsets all the way?  That is for
            // cdr and car? Should we have a utility function? All in all, we
            // need to refactor the codebase first to work with pointer offsets
            // rather than objects. This issue is very important, if we are in
            // the middle of an operation that tries to allocate memory but we
            // had saved pointers to some object, the pointer references may be
            // invalidated, crashing or worse, silently returning garbage! Let's
            // move on for now implementing the GC and we will revisit this part
            // later.
        }
    }
    size_t slot = gc.free_objects.offsets[gc.free_objects.position++];
    array_head(gc.free_objects.offsets)->size--;
    Object *obj = &gc.objects[slot];
    obj->type = type;
    obj->marked = false;
    return obj;
}