#include "suicmez_gc.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

typedef struct {
    /* Young generation */
    uint8_t *from_space;
    uint8_t *from_ptr;

    uint8_t *to_space;
    uint8_t *to_ptr;

    size_t young_size;

    /* Old generation */
    uint8_t *old;
    uint8_t *old_ptr;
    size_t old_size;

    /* Card table */
    uint8_t *card_table;
    size_t card_count;

    /* Roots = pointer slots */
    void **roots[MAX_ROOTS];
    size_t root_count;
} SuicmezGC;

static SuicmezGC gc;

static int is_young(void *p) {
    return (uint8_t*)p >= gc.from_space &&
           (uint8_t*)p <  gc.from_ptr;
}

static int is_old(void *p) {
    return (uint8_t*)p >= gc.old &&
           (uint8_t*)p <  gc.old_ptr;
}

static size_t card_index(void *p) {
    return ((uint8_t*)p - gc.old) / CARD_SIZE;
}

void gc_add_root(void **slot) {
    gc.roots[gc.root_count++] = slot;
}

void gc_remove_root(void **slot) {
    for (size_t i = 0; i < gc.root_count; i++) {
        if (gc.roots[i] == slot) {
            gc.roots[i] = gc.roots[--gc.root_count];
            return;
        }
    }
}

void gc_write_barrier(void *owner, void *value) {
    if (is_old(owner) && is_young(value)) {
        size_t idx = card_index(owner);
        if (idx < gc.card_count){
            gc.card_table[idx] = 1;
        }
    }
}

static Obj *promote(Obj *obj) {
    size_t sz = obj->header.size;
    if (gc.old_ptr + sz > gc.old + gc.old_size) {
        fprintf(stderr, "GC: old gen overflow\n");
        abort();
    }

    Obj *new_obj = (Obj*)gc.old_ptr;
    memcpy(new_obj, obj, sz);
    gc.old_ptr += sz;

    return new_obj;
}

static Obj *forward(Obj *obj) {
    if (!obj) return NULL;

    if (is_old(obj))
        return obj;

    if (!is_young(obj))
        return obj;

    if (obj->header.forwarding)
        return obj->header.forwarding;

    if (obj->header.age >= PROMOTION_AGE) {
        Obj *p = promote(obj);
        obj->header.forwarding = p;
        scan_object_fields(p);
        return p;
    }

    size_t sz = obj->header.size;
    Obj *new_obj = (Obj*)gc.to_ptr;
    memcpy(new_obj, obj, sz);

    gc.to_ptr += sz;

    new_obj->header.age++;
    new_obj->header.forwarding = NULL;
    obj->header.forwarding = new_obj;

    return new_obj;
}

static inline void scan_object_fields(Obj *obj) {
    const TypeInfo *t = obj->header.type;

    for (uint16_t i = 0; i < t->field_count; i++) {
        if (t->pointer_bitmap[i]) {
            obj->fields[i] = forward((Obj*)obj->fields[i]);
        }
    }
}

static void gc_minor_collect(void) {
    gc.to_ptr = gc.to_space;

    /* 1. Roots */
    for (size_t i = 0; i < gc.root_count; i++) {
        void **slot = gc.roots[i];
        *slot = forward((Obj*)*slot);
    }

    /* 2. Remembered set */
    for (size_t c = 0; c < gc.card_count; c++) {
        if (!gc.card_table[c]) continue;
        gc.card_table[c] = 0;

        uint8_t *start = gc.old + c * CARD_SIZE;
        uint8_t *end   = start + CARD_SIZE;
        if (end > gc.old_ptr) end = gc.old_ptr;

        for (uint8_t *p = start; p < end;) {
            Obj *obj = (Obj*)p;
            size_t fields =
                (obj->header.size - sizeof(ObjHeader)) / sizeof(void*);

            scan_object_fields(obj);
            p += obj->header.size;
        }
    }

    /* 3. Cheney scan */
    uint8_t *scan = gc.to_space;
    while (scan < gc.to_ptr) {
        Obj *obj = (Obj*)scan;
        size_t fields =
            (obj->header.size - sizeof(ObjHeader)) / sizeof(void*);

        scan_object_fields(obj);

        scan += obj->header.size;
    }

    /* 4. Swap spaces */
    uint8_t *tmp = gc.from_space;
    gc.from_space = gc.to_space;
    gc.to_space = tmp;
    gc.from_ptr = gc.to_ptr;

    /* 5. Clear forwarding pointers */
    scan = gc.from_space;
    while (scan < gc.from_ptr) {
        Obj *obj = (Obj*)scan;
        obj->header.forwarding = NULL;
        scan += obj->header.size;
    }
}

void *gc_alloc(const TypeInfo *type, size_t payload_size) {
    size_t size = sizeof(ObjHeader) + payload_size;
    size = (size + sizeof(void*) - 1) & ~(sizeof(void*) - 1);

    if (gc.from_ptr + size > gc.from_space + gc.young_size) {
        gc_minor_collect();
        if ((gc.old_ptr - gc.old) > (gc.old_size * 7 / 10)) {
            gc_collect_old();
        }
    }

    Obj *obj = (Obj*)gc.from_ptr;
    obj->header.size = size;
    obj->header.age = 0;
    obj->header.flags = 0;
    obj->header.forwarding = NULL;
    obj->header.type = type;

    gc.from_ptr += size;
    return obj;
}

void gc_init(void) {
    gc.young_size = YOUNG_SIZE;
    gc.old_size   = OLD_SIZE;

    gc.from_space = malloc(gc.young_size);
    gc.to_space   = malloc(gc.young_size);
    gc.old        = malloc(gc.old_size);

    gc.from_ptr = gc.from_space;
    gc.to_ptr   = gc.to_space;
    gc.old_ptr  = gc.old;

    gc.card_count = gc.old_size / CARD_SIZE;
    gc.card_table = calloc(gc.card_count, 1);

    gc.root_count = 0;
}

static void mark_object_fields(Obj *obj) {
    const TypeInfo *t = obj->header.type;
    for (uint16_t i = 0; i < t->field_count; i++) {
        if (t->pointer_bitmap[i]) {
            Obj *child = (Obj*)obj->fields[i];
            if (child && is_old(child)) {
                mark_obj(child);
            }
        }
    }
}

static void mark_obj(Obj *obj) {
    if (!obj) return;
    if (!(obj->header.flags & MARKED)) {
        obj->header.flags |= MARKED;
        mark_object_fields(obj);
    }
}

static void mark_roots(void) {
    for (size_t i = 0; i < gc.root_count; i++) {
        Obj *obj = (Obj*)*gc.roots[i];
        if (obj && is_old(obj)) {
            mark_obj(obj);
        }
    }
}

static void compute_forwarding(void) {
    uint8_t *dest = gc.old;
    uint8_t *scan = gc.old;

    while (scan < gc.old_ptr) {
        Obj *obj = (Obj*)scan;
        if (obj->header.flags & MARKED) {
            obj->header.forwarding = (Obj*)dest;
            dest += obj->header.size;
        }
        scan += obj->header.size;
    }
}

static void update_old_pointers(void) {
    uint8_t *scan = gc.old;

    while (scan < gc.old_ptr) {
        Obj *obj = (Obj*)scan;
        if (obj->header.flags & MARKED) {
            const TypeInfo *t = obj->header.type;
            for (uint16_t i = 0; i < t->field_count; i++) {
                if (t->pointer_bitmap[i]) {
                    Obj *child = (Obj*)obj->fields[i];
                    if (child && is_old(child)) {
                        obj->fields[i] = child->header.forwarding;
                    }
                }
            }
        }
        scan += obj->header.size;
    }
}

static void update_root_pointers(void) {
    for (size_t i = 0; i < gc.root_count; i++) {
        Obj *obj = (Obj*)*gc.roots[i];
        if (obj && is_old(obj)) {
            *gc.roots[i] = obj->header.forwarding;
        }
    }
}

static void compact_old(void) {
    uint8_t *scan = gc.old;
    uint8_t *dest = gc.old;

    while (scan < gc.old_ptr) {
        Obj *obj = (Obj*)scan;
        if (obj->header.flags & MARKED) {
            Obj *new_loc = obj->header.forwarding;
            memmove(new_loc, obj, obj->header.size);

            Obj *moved = new_loc;
            moved->header.flags &= ~MARKED;
            moved->header.forwarding = NULL;

            dest += moved->header.size;
        }
        scan += obj->header.size;
    }

    gc.old_ptr = dest;
}

void gc_collect_old(void) {
    mark_roots();
    compute_forwarding();
    update_root_pointers();
    update_old_pointers();
    compact_old();
}

void gc_shutdown(void) {
    free(gc.from_space);
    free(gc.to_space);
    free(gc.old);
    free(gc.card_table);
}