#include "arena.hpp"

#include "die.hpp"
#include "error.hpp"
#include "pod.hpp"

static Arena* current_arena = 0;

GcRootBase::GcRootBase(PodObject* ptr, GcRootList* node)
    : _ptr(ptr), _node(node) {}

GcRootBase::~GcRootBase() {
  if (_node != 0) _node->remove();
}

Result<void> Arena::gc() {
  std::swap(_first, _second);
  _first->reset();

  TRY(gc_roots());

  _second->reset();
  return Result<void>();
}

Result<void> Arena::gc_roots() {
  auto node = _gcroot.next();
  while (node != 0) {
    node = TRY(gc_root(node));
    node = node->next();
  }
  if (_gc_hold != 0) {
    _gc_hold = TRY(gc_pod(_gc_hold));
  }
  return Result<void>();
}

Result<GcRootList*> Arena::gc_root(GcRootList* node) {
  auto pod = TRY(gc_pod(node->_root->_ptr));
  auto lst = TRY(alloc_raw<GcRootList>());
  node->_root->_ptr = pod;
  node->_root->_node = lst;

  lst->_root = node->_root;

  GcRootList* prev = node->_prev;
  GcRootList* next = node->_next;

  prev->_next = lst;
  if (next) next->_prev = lst;
  lst->_prev = prev;
  lst->_next = next;

  // Not strictly needed, but useful for debugging
  node->_prev = 0;
  node->_next = 0;
  node->_root = 0;

  return lst;
}

Result<PodObject*> Arena::gc_pod(PodObject* obj) {
  if (obj == 0) return 0;

  switch (obj->header.tag) {
    case Tag::Nil:
      return gc_nil((PodNil*)obj);
    case Tag::Bool:
      return gc_bool((PodBool*)obj);
    case Tag::Int64:
      return gc_int64((PodInt64*)obj);
    case Tag::Float:
      return gc_float((PodFloat*)obj);
    case Tag::String:
      return gc_string((PodString*)obj);
    case Tag::Symbol:
      return gc_symbol((PodSymbol*)obj);
    case Tag::SrcLoc:
      return gc_srcloc((PodSrcLoc*)obj);
    case Tag::Syntax:
      return gc_syntax((PodSyntax*)obj);
    case Tag::Pair:
      return gc_pair((PodPair*)obj);
    case Tag::Array:
      return gc_array((PodArray*)obj);
    case Tag::ByteArray:
      return gc_bytearray((PodByteArray*)obj);
    case Tag::Dict:
      return gc_dict((PodDict*)obj);
    case Tag::Opcode:
      return gc_opcode((PodOpcode*)obj);
    case Tag::Function:
      return gc_function((PodFunction*)obj);
    case Tag::StdlibFunction:
      return gc_stdlib_function((PodStdlibFunction*)obj);
    case Tag::Module:
      return gc_module((PodModule*)obj);
    case Tag::StackFrame:
      return gc_stack((PodStackFrame*)obj);
    case Tag::Error:
      return gc_error((PodError*)obj);
    case Tag::Continuation:
      return gc_continuation((PodContinuation*)obj);
    case Tag::Task:
      return gc_task((PodTask*)obj);
    case Tag::TaskResult:
      return gc_task_result((PodTaskResult*)obj);
  }

  return ERROR(TypeMismatch);
}

Result<PodObject*> Arena::gc_nil(PodNil* obj) {
  auto nobj = TRY(alloc_raw<PodNil>());
  nobj->header.tag = Tag::Nil;
  return nobj;
}

Result<PodObject*> Arena::gc_bool(PodBool* obj) {
  auto nobj = TRY(alloc_raw<PodBool>());
  nobj->header.tag = Tag::Bool;
  nobj->value = obj->value;
  return nobj;
}

Result<PodObject*> Arena::gc_int64(PodInt64* obj) {
  auto nobj = TRY(alloc_raw<PodInt64>());
  nobj->header.tag = Tag::Int64;
  nobj->value = obj->value;
  return nobj;
}

Result<PodObject*> Arena::gc_float(PodFloat* obj) {
  auto nobj = TRY(alloc_raw<PodFloat>());
  nobj->header.tag = Tag::Float;
  nobj->value = obj->value;
  return nobj;
}

Result<PodObject*> Arena::gc_string(PodString* obj) {
  auto nobj = TRY(alloc_raw<PodString>(obj->size * sizeof(char32_t)));
  nobj->header.tag = Tag::String;
  nobj->size = obj->size;
  memcpy(nobj->data, obj->data, obj->size * sizeof(char32_t));
  return nobj;
}

Result<PodObject*> Arena::gc_symbol(PodSymbol* obj) {
  auto nobj = TRY(alloc_raw<PodSymbol>(obj->size * sizeof(char32_t)));
  nobj->header.tag = Tag::Symbol;
  nobj->size = obj->size;
  memcpy(nobj->data, obj->data, obj->size * sizeof(char32_t));
  return nobj;
}

Result<PodObject*> Arena::gc_srcloc(PodSrcLoc* obj) {
  auto nobj = TRY(alloc_raw<PodSrcLoc>());

  nobj->header.tag = Tag::SrcLoc;

  nobj->sourcerange = obj->sourcerange;

  return nobj;
}

Result<PodObject*> Arena::gc_syntax(PodSyntax* obj) {
  auto nobj = TRY(alloc_raw<PodSyntax>());

  nobj->header.tag = Tag::Syntax;

  nobj->filename = TRY(gc_pod(obj->filename.get()));
  nobj->srcloc = TRY(gc_pod(obj->srcloc.get()));
  nobj->expression = TRY(gc_pod(obj->expression.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_pair(PodPair* obj) {
  auto nobj = TRY(alloc_raw<PodPair>());

  nobj->header.tag = Tag::Pair;

  PodObject* first = obj->first.get();
  PodObject* rest = obj->rest.get();

  first = TRY(gc_pod(first));
  rest = TRY(gc_pod(rest));

  nobj->first = first;
  nobj->rest = rest;

  return nobj;
}

Result<PodObject*> Arena::gc_array(PodArray* obj) {
  auto nobj = TRY(alloc_raw<PodArray>(sizeof(OffPtr<PodObject>) * obj->size));
  nobj->header.tag = Tag::Array;
  nobj->size = obj->size;
  for (uint64_t i = 0; i < obj->size; i++) {
    PodObject* val = obj->data[i].get();
    nobj->data[i] = TRY(gc_pod(val));
  }
  return nobj;
}

Result<PodObject*> Arena::gc_bytearray(PodByteArray* obj) {
  auto nobj = TRY(alloc_raw<PodByteArray>(obj->size));
  nobj->header.tag = Tag::ByteArray;
  nobj->size = obj->size;
  memcpy(nobj->data, obj->data, obj->size);
  return nobj;
}

Result<PodObject*> Arena::gc_dict(PodDict* obj) {
  auto nobj =
      TRY(alloc_raw<PodDict>(obj->size * 2 * sizeof(OffPtr<PodObject>)));
  nobj->header.tag = Tag::Dict;
  nobj->size = obj->size;
  for (uint64_t i = 0; i < obj->size * 2; i++) {
    PodObject* val = obj->data[i].get();
    nobj->data[i] = TRY(gc_pod(val));
  }
  return nobj;
}

Result<PodObject*> Arena::gc_opcode(PodOpcode* obj) {
  auto nobj = TRY(alloc_raw<PodOpcode>());
  nobj->header.tag = Tag::Opcode;
  nobj->opcode = obj->opcode;
  nobj->arg1 = obj->arg1;
  nobj->arg2 = obj->arg2;
  nobj->arg3 = obj->arg3;
  nobj->arg4 = obj->arg4;

  return nobj;
}

Result<PodObject*> Arena::gc_function(PodFunction* obj) {
  auto nobj = TRY(alloc_raw<PodFunction>());
  nobj->header.tag = Tag::Function;
  nobj->arity = obj->arity;
  nobj->name = TRY(gc_pod(obj->name.get()));
  nobj->constants = TRY(gc_pod(obj->constants.get()));
  nobj->code = TRY(gc_pod(obj->code.get()));
  nobj->closure = TRY(gc_pod(obj->closure.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_stdlib_function(PodStdlibFunction* obj) {
  auto nobj = TRY(alloc_raw<PodStdlibFunction>());
  nobj->header.tag = Tag::StdlibFunction;
  nobj->fun_id = obj->fun_id;

  return nobj;
}

Result<PodObject*> Arena::gc_module(PodModule* obj) {
  auto nobj = TRY(alloc_raw<PodModule>());
  nobj->header.tag = Tag::Module;
  nobj->name = TRY(gc_pod(obj->name.get()));
  nobj->fun = TRY(gc_pod(obj->fun.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_stack(PodStackFrame* obj) {
  auto nobj =
      TRY(alloc_raw<PodStackFrame>(sizeof(OffPtr<PodObject>) * obj->size));
  nobj->header.tag = Tag::StackFrame;
  nobj->parent = TRY(gc_pod(obj->parent.get()));
  nobj->fun = TRY(gc_pod(obj->fun.get()));
  nobj->guard = obj->guard;
  nobj->pc = obj->pc;
  nobj->size = obj->size;
  for (uint64_t i = 0; i < obj->size; i++) {
    PodObject* val = obj->data[i].get();
    nobj->data[i] = TRY(gc_pod(val));
  }
  return nobj;
}

Result<PodObject*> Arena::gc_error(PodError* obj) {
  auto nobj = TRY(alloc_raw<PodError>());
  nobj->header.tag = Tag::Error;

  nobj->name = TRY(gc_pod(obj->name.get()));
  nobj->message = TRY(gc_pod(obj->message.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_continuation(PodContinuation* obj) {
  auto nobj = TRY(alloc_raw<PodContinuation>());
  nobj->header.tag = Tag::Continuation;

  nobj->value = TRY(gc_pod(obj->value.get()));
  nobj->frame = TRY(gc_pod(obj->frame.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_task(PodTask* obj) {
  auto nobj = TRY(alloc_raw<PodTask>());
  nobj->header.tag = Tag::Task;

  nobj->task_id = obj->task_id;
  nobj->params = TRY(gc_pod(obj->params.get()));

  return nobj;
}

Result<PodObject*> Arena::gc_task_result(PodTaskResult* obj) {
  auto nobj = TRY(alloc_raw<PodTaskResult>());
  nobj->header.tag = Tag::TaskResult;

  nobj->result = TRY(gc_pod(obj->result.get()));
  nobj->error = TRY(gc_pod(obj->error.get()));

  return nobj;
}

Arena& get_arena() {
  if (current_arena == 0) die("Arena not set\n");
  return *current_arena;
}
void set_arena(Arena& arena) { current_arena = &arena; }

Result<void> arena_gc() { return get_arena().gc(); }