#include "gc_mark.h"
#include "gc_roots.h"
#include "../runtime.h"
#include <string.h>

static RavaGCState_t g_gc_state;

void rava_gc_mark_init(void) {
    memset(&g_gc_state, 0, sizeof(g_gc_state));
    g_gc_state.gc_enabled = true;
}

RavaGCState_t* rava_gc_get_state(void) {
    return &g_gc_state;
}

void rava_gc_mark_object(RavaGCHeader_t *obj) {
    if (!obj || obj->gc_mark != RAVA_GC_WHITE) {
        return;
    }

    obj->gc_mark = RAVA_GC_GRAY;

    if (g_gc_state.gray_count < RAVA_GC_GRAY_STACK_SIZE) {
        g_gc_state.gray_stack[g_gc_state.gray_count++] = obj;
    } else {
        rava_gc_process_object(obj);
    }
}

static void _rava_gc_trace_object(RavaObject_t *obj) {
    for (size_t i = 0; i < obj->field_count; i++) {
        rava_gc_mark_value(obj->field_values[i]);
    }
}

static void _rava_gc_trace_array(RavaArray_t *arr) {
    if (arr->element_type == RAVA_VAL_OBJECT ||
        arr->element_type == RAVA_VAL_ARRAY) {
        RavaValue_t *values = (RavaValue_t*)arr->data;
        for (size_t i = 0; i < arr->length; i++) {
            rava_gc_mark_value(values[i]);
        }
    }
}

static void _rava_gc_trace_arraylist(RavaArrayList_t *list) {
    for (size_t i = 0; i < list->size; i++) {
        rava_gc_mark_value(list->data[i]);
    }
}

static void _rava_gc_trace_hashmap(RavaHashMap_t *map) {
    for (size_t i = 0; i < map->bucket_count; i++) {
        if (map->buckets[i].occupied) {
            rava_gc_mark_value(map->buckets[i].value);
        }
    }
}

static void _rava_gc_trace_methodref(RavaMethodRef_t *ref) {
    if (ref->has_target) {
        rava_gc_mark_value(ref->target);
    }
}

void rava_gc_process_object(RavaGCHeader_t *obj) {
    obj->gc_mark = RAVA_GC_BLACK;
    g_gc_state.objects_marked++;
    g_gc_state.bytes_marked += obj->gc_size;

    switch (obj->gc_type) {
        case RAVA_GC_TYPE_OBJECT: {
            RavaObject_t *o = (RavaObject_t*)obj;
            _rava_gc_trace_object(o);
            break;
        }
        case RAVA_GC_TYPE_ARRAY: {
            RavaArray_t *a = (RavaArray_t*)obj;
            _rava_gc_trace_array(a);
            break;
        }
        case RAVA_GC_TYPE_ARRAYLIST: {
            RavaArrayList_t *l = (RavaArrayList_t*)obj;
            _rava_gc_trace_arraylist(l);
            break;
        }
        case RAVA_GC_TYPE_HASHMAP: {
            RavaHashMap_t *m = (RavaHashMap_t*)obj;
            _rava_gc_trace_hashmap(m);
            break;
        }
        case RAVA_GC_TYPE_METHODREF: {
            RavaMethodRef_t *r = (RavaMethodRef_t*)obj;
            _rava_gc_trace_methodref(r);
            break;
        }
        default:
            break;
    }
}

void rava_gc_process_gray_stack(void) {
    while (g_gc_state.gray_count > 0) {
        RavaGCHeader_t *obj = g_gc_state.gray_stack[--g_gc_state.gray_count];
        rava_gc_process_object(obj);
    }
}