#include "memory.h"
#include "vm.h"

#ifdef DEBUG_LOG_GC
#include "debug.h"
#include <stdlib.h>
#endif

#define GC_HEAP_GROW_FACTOR 2

static void markArray(ValueArray *);

//    oldSize	 newSize	               Operation
//    0	        Non‑zero	            Allocate new block.
//    Non‑zero	0	                    Free allocation.
//    Non‑zero	Smaller than oldSize	Shrink existing allocation.
//    Non‑zero	Larger than oldSize	    Grow existing allocation.
void *reallocate(void *pointer, size_t oldSize, size_t newSize) {
  vm.bytesAllocated += newSize - oldSize;
  if (newSize > oldSize) {
#ifdef DEBUG_STRESS_GC
    collectGarbage();
#endif

    if (vm.bytesAllocated > vm.nextGC) {
      collectGarbage();
    }
  }

  if (newSize == 0) {
    free(pointer);
    return NULL;
  }

  /*
  The realloc() function tries to change the size of the allocation pointed
   to by ptr to size, and returns ptr.  If there is not enough room to
   enlarge the memory allocation pointed to by ptr, realloc() creates a new
   allocation, copies as much of the old data pointed to by ptr as will fit
   to the new allocation, frees the old allocation, and returns a pointer to
   the allocated memory.  If ptr is NULL, realloc() is identical to a call
   to malloc() for size bytes.  If size is zero and ptr is not NULL, a new,
   minimum sized object is allocated and the original object is freed.  When
   extending a region allocated with calloc(3), realloc(3) does not guaran-
   tee that the additional memory is also zero-filled
   */
  void *result = realloc(pointer, newSize);
  if (result == NULL) exit(1);
  return result;
}
static void freeObject(Obj *object) {
#ifdef DEBUG_LOG_GC
  PRINTLNF("%p free type %d\n", (void *) object, object->type);
#endif

  switch (object->type) {
    case OBJ_NATIVE:
      FREE(ObjNative, object);
      break;
    case OBJ_FUNCTION: {
      ObjFunction *function = (ObjFunction *) object;
      freeChunk(&function->chunk);
      FREE(ObjFunction, object);
      break;
    }
    case OBJ_STRING:
      FREE_ARRAY(char, ((ObjString *) object)->chars, ((ObjString *) object)->length + 1);
      FREE(ObjString, object);
      break;
    case OBJ_CLOSURE: {
      ObjClosure *closure = (ObjClosure *) object;
      FREE_ARRAY(ObjUpvalue *, closure->upvalues, closure->upvalueCount);
      FREE(ObjClosure, object);
      break;
    }
    case OBJ_UPVALUE:
      FREE(ObjUpvalue, object);
      break;
    case OBJ_CLASS: {
      ObjClass *klass = (ObjClass *) object;
      freeTable(&klass->methods);
      FREE(ObjClass, object);
      break;
    }
    case OBJ_INSTANCE: {
      ObjInstance *instance = (ObjInstance *) object;
      freeTable(&instance->fields);
      FREE(ObjInstance, object);
      break;
    }
    case OBJ_BOUND_METHOD:
      FREE(ObjBoundMethod, object);
      break;
  }
}
void freeObjects() {
  Obj *object = vm.objects;
  while (object != NULL) {
    Obj *next = object->next;
    freeObject(object);
    object = next;
  }

  free(vm.grayStack);
}
static void markRoots() {
  for (Value *slot = vm.stack; slot < vm.stackTop; slot++) {
    markValue(*slot);
  }

  for (int i = 0; i < vm.frameCount; i++) {
    markObject((Obj *) vm.frames[i].closure);
  }

  for (ObjUpvalue *upvalue = vm.openUpvalues;
       upvalue != NULL; upvalue = upvalue->next) {
    markObject((Obj *) upvalue);
  }

  markTable(&vm.globals);
  markCompilerRoots();
}
void blackenObject(Obj *object) {
#ifdef DEBUG_LOG_GC
  printf("%p blacken ", (void *) object);
  printValue(OBJ_VAL(object));
  printf("\n");
#endif

  switch (object->type) {
    case OBJ_INSTANCE: {
      ObjInstance *instance = (ObjInstance *) object;
      markObject((Obj *) instance->klass);
      markTable(&instance->fields);
      break;
    }
    case OBJ_CLASS: {
      ObjClass *klass = (ObjClass *) object;
      markObject((Obj *) klass->name);
      markTable(&klass->methods);
      break;
    }
    case OBJ_CLOSURE: {
      ObjClosure *closure = (ObjClosure *) object;
      markObject((Obj *) closure->function);
      for (int i = 0; i < closure->upvalueCount; i++) {
        markObject((Obj *) closure->upvalues[i]);
      }
      break;
    }
    case OBJ_FUNCTION: {
      ObjFunction *function = (ObjFunction *) object;
      markObject((Obj *) function->name);
      markArray(&function->chunk.constants);
      break;
    }
    case OBJ_UPVALUE:
      markValue(((ObjUpvalue *) object)->closed);
      break;
    case OBJ_NATIVE:
    case OBJ_STRING:
      break;
    case OBJ_BOUND_METHOD: {
      ObjBoundMethod *bound = (ObjBoundMethod *) object;
      markValue(bound->receiver);
      markObject((Obj *) bound->method);
      break;
    }
  }
}
void traceReferences() {
  while (vm.grayCount > 0) {
    Obj *object = vm.grayStack[--vm.grayCount];
    blackenObject(object);
  }
}
void sweep() {
  Obj *previous = NULL;
  Obj *object = vm.objects;
  while (object != NULL) {
    if (object->isMarked) {
      object->isMarked = false;
      previous = object;
      object = object->next;
    } else {
      Obj *unreached = object;
      object = object->next;
      if (previous != NULL) {
        previous->next = object;
      } else {
        vm.objects = object;
      }

      freeObject(unreached);
    }
  }
}
void collectGarbage() {
#ifdef DEBUG_LOG_GC
  printf("-- gc begin\n");
  size_t before = vm.bytesAllocated;
#endif

  markRoots();
  traceReferences();
  tableRemoveWhite(&vm.strings);
  sweep();// sweeping unused objects

  vm.nextGC = vm.bytesAllocated * GC_HEAP_GROW_FACTOR;

#ifdef DEBUG_LOG_GC
  printf("-- gc end\n");
  printf("   collected %zu bytes (from %zu to %zu) next at %zu\n",
         before - vm.bytesAllocated, before, vm.bytesAllocated, vm.nextGC);
#endif
}
void markValue(Value value) {
  if (IS_OBJ(value)) markObject(AS_OBJ(value));
}
static void markArray(ValueArray *array) {
  for (int i = 0; i < array->count; i++) {
    markValue(array->values[i]);
  }
}
void markObject(Obj *object) {
  if (object == NULL) return;
  if (object->isMarked) return;
#ifdef DEBUG_LOG_GC
  printf("%p mark ", (void *) object);
  printValue(OBJ_VAL(object));
  printf("\n");
#endif

  object->isMarked = true;

  if (vm.grayCapacity < vm.grayCount + 1) {
    vm.grayCapacity = GROW_CAPACITY(vm.grayCapacity);
    vm.grayStack = (Obj **) realloc(vm.grayStack, sizeof(Obj *) * vm.grayCapacity);

    if (vm.grayStack == NULL) exit(1);
  }

  vm.grayStack[vm.grayCount++] = object;
}