#include "gc.h"

Environment *
alloc_env(void) {
    if (gc.free_envs.size == 0) {
        mark_and_sweep();
        if (gc.free_envs.size == 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.buf[gc.free_envs.position++];
    gc.free_envs.size--;
    return &gc.envs.buf[slot];
}

void
push_root(Object *obj) {
    if (gc.roots.size == gc.roots.cap) {
        gc.roots.cap *= 2;
        gc.roots.buf = realloc(gc.roots.buf, gc.roots.cap * sizeof(Object *));
    }
    gc.roots.buf[gc.roots.size++] = obj;
}

Object *
pop_root(void) {
    return gc.roots.buf[gc.roots.size--];
}

void
push_active_env(Environment *env) {
    if (gc.active_envs.size == gc.active_envs.cap) {
        gc.active_envs.cap *= 2;
        gc.active_envs.buf = realloc(gc.active_envs.buf,
                gc.active_envs.cap * sizeof(Environment *));
    }
    gc.active_envs.buf[gc.active_envs.size++] = env;
}

Environment *
pop_active_env(void) {
    return gc.active_envs.buf[gc.active_envs.size--];
}

void
init_gc(void) {
    gc = (GC){
        .objects = malloc(GC_OBJS_CAP * sizeof(Object)),
        .obj_cap = GC_OBJS_CAP,
        .envs = (Environments){
            .buf = malloc(GC_ENVS_CAP * sizeof(Environment)),
            .size = 0,
            .cap = GC_ENVS_CAP,
        },
        .roots = (RootNodes){
            .buf = malloc(GC_ROOTS_CAP * sizeof(Object*)),
            .size = 0,
            .cap = GC_ROOTS_CAP,
        },
        .free_objects = (FreeList){
            .buf = malloc(GC_OBJS_CAP * sizeof(size_t)),
            .size = GC_OBJS_CAP,
            .cap = GC_OBJS_CAP,
        },
        .free_envs = (FreeList){
            .buf = malloc(GC_ENVS_CAP * sizeof(size_t)),
            .size = GC_ENVS_CAP,
            .cap = GC_ENVS_CAP,
        },
        .active_envs = (ActiveEnvs){
            .buf = malloc(GC_ACTIVE_ENVS_CAP * sizeof(Environment*)),
            .size = 0,
            .cap = GC_ACTIVE_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++) {
        gc.free_objects.buf[i] = i;
    }
    for (size_t i = 0; i < GC_ENVS_CAP; i++) {
        gc.free_envs.buf[i] = i;
    }
}

void
mark_environment(Environment *env) {
    if (env == NULL || env->marked) {
        return;
    }
    env->marked = true;
    for (size_t i = 0; i < env->size; i++) {
        EnvEntry *entry = &env->buf[i];
        mark_obj(entry->symbol);
        mark_obj(entry->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 < gc.active_envs.size; i++) {
        mark_environment(gc.active_envs.buf[i]);
    }
    for (size_t i = 0; i < gc.roots.size; i++) {
        mark_obj(gc.roots.buf[i]);
    }

    // Reset the free list.
    gc.free_objects.position = 0;
    gc.free_objects.size = 0;
    gc.free_envs.position = 0;
    gc.free_envs.size = 0;

    // Sweep.
    for (size_t i = 0; i < gc.obj_cap; i++) {
        Object *obj = &gc.objects[i];
        if (!obj->marked) {
            // Free heap allocated memory for this object if needed.
            if (obj->type == OBJ_TYPE_SYMBOL) {
                if (obj->symbol != NULL) {
                    free(obj->symbol);
                }
                obj->symbol = NULL;
                obj->symbol_n = 0;
            } else if (obj->type == OBJ_TYPE_STRING) {
                if (obj->symbol != NULL) {
                    free(obj->string);
                }
                obj->string = NULL;
                obj->string_n = 0;
            }
            gc.free_objects.buf[gc.free_objects.size++] = i;
        }
        obj->marked = false;
    }
    for (size_t i = 0; i < gc.envs.cap; i++) {
        Environment *env = &gc.envs.buf[i];
        if (!env->marked) {
            if (env->buf != NULL) {
                free(env->buf);
                env->buf = NULL;
                env->size = 0;
                env->cap = 0;
            }
            gc.free_envs.buf[gc.free_envs.size++] = i;
        }
        env->marked = false;
    }
}

void
dump_gc(void) {
        printf("-------------- ROOTS -------------- \n");
        for (size_t i = 0; i < gc.roots.size; i++) {
            display(gc.roots.buf[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 < gc.obj_cap; j++) {
                if (gc.free_objects.buf[j] == i) {
                    is_free = true;
                    break;
                }
            }
            if (is_free) {
                printf(" [FREE]");
            }
            printf("\n");
        }
        printf("-------------- MISC --------------- \n");
        printf("gc.roots.size: %ld\n", gc.roots.size);
        printf("gc.roots.cap: %ld\n", gc.roots.cap);
        printf("gc.active_envs.size: %ld\n", gc.active_envs.size);
        printf("gc.active_envs.cap: %ld\n", gc.active_envs.cap);
        printf("gc.obj_cap: %ld\n", gc.obj_cap);
        printf("gc.free_objects.size: %ld\n", gc.free_objects.size);
        printf("gc.free_objects.cap: %ld\n", gc.free_objects.cap);
        printf("gc.free_objects.position: %ld\n", gc.free_objects.position);
        printf("gc.free_envs.size: %ld\n", gc.free_envs.size);
        printf("gc.free_envs.cap: %ld\n", gc.free_envs.cap);
        printf("gc.free_envs.position: %ld\n", gc.free_envs.position);
        printf("gc.envs.size: %ld\n", gc.envs.size);
        printf("gc.envs.cap: %ld\n", gc.envs.cap);
}

Object *
alloc_object(ObjectType type) {
    if (gc.free_objects.size == 0) {
        mark_and_sweep();
        if (gc.free_objects.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.buf[gc.free_objects.position++];
    gc.free_objects.size--;
    Object *obj = &gc.objects[slot];
    obj->type = type;
    obj->marked = false;
    return obj;
}