// Stack of root nodes.
typedef struct RootNodes {
    Object **buf;
    size_t size;
    size_t cap;
} RootNodes;

typedef struct Environments {
    Environment *buf;
    size_t size;
    size_t cap;
} Environments;

typedef struct GC {
    RootNodes roots;
    Environments envs;
    Object *obj_list;
    // Free list keeps track of the offset numbers from the obj_list
    size_t *free_list;
    size_t obj_cap;
    size_t fl_pos;
    size_t available_slots;
} GC;

// FIXME: small value for testing purposes
// #define GC_INITIAL_HEAP 32
#define GC_INITIAL_HEAP 1024 * 1.5
#define GC_ROOTS_CAP 1024 * 1024
#define GC_ENVS_CAP 1024 * 1024

static GC gc;

Environment *
alloc_env(void) {
    if (gc.envs.size < gc.envs.cap) {
        return &gc.envs.buf[gc.envs.size++];
    }
    printf("error: not enough room for more environments\n");
    return NULL;
}

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
init_gc(void) {
    gc = (GC){
        .free_list = malloc(GC_INITIAL_HEAP * sizeof(size_t)),
        .obj_list = malloc(GC_INITIAL_HEAP * sizeof(Object)),
        .obj_cap = GC_INITIAL_HEAP,
        .available_slots = GC_INITIAL_HEAP,
        .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,
        },
    };

    // The free list stores the offset from the initial position for all
    // available slots.
    for (size_t i = 0; i < gc.obj_cap; i++) {
        gc.free_list[i] = i;
    }
}

Object *
get_obj(size_t offset) {
    return &gc.obj_list[offset];
}

void mark_obj(Object *obj);

void
mark_environment(Environment *env) {
    if (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.envs.size; i++) {
        mark_environment(&gc.envs.buf[i]);
    }

    for (size_t i = 0; i < gc.roots.size; i++) {
        if (gc.roots.buf[i]->marked) {
            continue;
        }
        mark_obj(gc.roots.buf[i]);
    }
    // dump_gc()

    // Reset the free list.
    gc.fl_pos = 0;
    gc.available_slots = 0;

    // Sweep.
    for (size_t i = 0; i < gc.obj_cap; i++) {
        Object *obj = &gc.obj_list[i];
        if (!obj->marked) {
            // Free heap allocated memory for this object if needed.
            if (obj->type == OBJ_TYPE_SYMBOL) {
                free(obj->symbol);
            } else if (obj->type == OBJ_TYPE_STRING) {
                free(obj->string);
            }
            gc.free_list[gc.available_slots++] = i;
        }
        obj->marked = false;
    }
    for (size_t i = 0; i < gc.envs.size; i++) {
        gc.envs.buf[i].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 ------------- \n");
        // for (size_t i = 0; i < gc.obj_cap; i++) {
        for (size_t i = 0; i < 20; i++) {
            printf("i: %ld -> ", i);
            Object *obj = &gc.obj_list[i];
            display(obj);
            bool is_free = false;
            for (size_t j = 0; j < gc.obj_cap; j++) {
                if (gc.free_list[j] == i) {
                    is_free = true;
                    break;
                }
            }
            if (is_free) {
                printf(" [FREE]");
            }
            printf("\n");
        }
        printf("FREE OBJECTS: %ld\n", gc.available_slots);
        printf("ENVIRONMENTS: %ld\n", gc.envs.size);
}

Object *
alloc_object(ObjectType type) {
    if (gc.available_slots == 0) {
        mark_and_sweep();
        if (gc.available_slots == 0) {
            printf("NOT MORE 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_list[gc.fl_pos++];
    gc.available_slots--;
    Object *obj = get_obj(slot);
    obj->type = type;
    return obj;
}