#pragma once

#include <sanitizer/asan_interface.h>

#include <concepts>
#include <cstring>
#include <iostream>

#include "error.hpp"
#include "pod.hpp"
#include "result.hpp"

class Arena;
class GcRootList;

class GcRootBase {
 public:
  GcRootBase() : _ptr(0), _node(0){};
  GcRootBase(PodObject* ptr, GcRootList* node);

  ~GcRootBase();

  PodObject* get() { return _ptr; }

 protected:
  PodObject* _ptr;
  GcRootList* _node;
};

template <class T>
  requires std::derived_from<T, PodObject>
class GcRoot : public GcRootBase {
 public:
  GcRoot() : GcRootBase(0, 0) {}
  GcRoot(T* ptr, GcRootList* node) : GcRootBase(ptr, node) {}
  GcRoot(GcRoot&& rhs);
  GcRoot& operator=(GcRoot&& rhs);

  static Result<GcRoot<T>> create(T* ptr, Arena& arena);
  Result<GcRoot<T>> copy(Arena& arena) {
    return GcRoot<T>::create((T*)_ptr, arena);
  }

  T* get() { return (T*)_ptr; }
  T& operator*() { return *(T*)_ptr; }
  T* operator->() { return (T*)_ptr; }
};

template <class T>
  requires std::derived_from<T, PodObject>
Result<GcRoot<T>> MkGcRoot(T* ptr, Arena& arena) {
  return GcRoot<T>::create(ptr, arena);
}

class GcRootList {
 public:
  GcRootList() : _prev(0), _next(0), _root(0) {}

  GcRootList(GcRootBase* root) : _prev(0), _next(0), _root(root) {}
  void insert(GcRootList* node) {
    node->_next = _next;
    node->_prev = this;
    if (_next) _next->_prev = node;
    _next = node;
  }

  void remove() {
    if (_prev) _prev->_next = _next;
    if (_next) _next->_prev = _prev;
    _prev = 0;
    _next = 0;
  }

  GcRootList* next() { return _next; }

  void update(GcRootBase* root) { _root = root; }

 private:
  GcRootList* _prev;
  GcRootList* _next;
  GcRootBase* _root;
};

class ArenaHeap {
 public:
  ArenaHeap(uint8_t* buf, uint64_t bufsize)
      : buf(buf), bufsize(bufsize), boundary(0) {
    memset(buf, 0, bufsize);
    ASAN_POISON_MEMORY_REGION(buf, bufsize);
  }

  Result<void*> alloc(uint64_t size) {
    // Always align to 64 bits
    if (size % 8 != 0) size += 8 - size % 8;

    // padding for guard value
    uint64_t boundary_size = 8;

    if (boundary + size + boundary_size >= bufsize)
      return ErrorCode::OutOfMemory;
    void* res = buf + boundary;
    boundary += size;
    ASAN_UNPOISON_MEMORY_REGION(res, size);
    boundary += boundary_size;
    return res;
  }

 private:
  uint8_t* buf;
  uint64_t bufsize;
  uint64_t boundary;
};

class Arena {
 public:
  Arena(ArenaHeap* first, ArenaHeap* second)
      : _heaps{first, second}, _gcroot(), _current(0) {}

  template <class T>
  Result<T*> alloc(uint64_t extra = 0) {
    uint64_t objsize = sizeof(T) + extra;
    auto ptr = _heaps[_current]->alloc(objsize);

    if (ptr.has_value()) return (T*)ptr.value();

    // TODO: trigger GC

    return (T*)TRY(_heaps[_current]->alloc(objsize));
  }

  void add_root(GcRootList* node) { _gcroot.insert(node); }

  uint64_t root_count() {
    uint64_t res = 0;
    for (GcRootList* cur = &_gcroot; cur->next() != 0; cur = cur->next()) {
      ++res;
    }
    return res;
  };

 private:
  ArenaHeap* _heaps[2];
  GcRootList _gcroot;
  int _current;
};

template <uint64_t size>
class StaticArenaHeap : public ArenaHeap {
 public:
  StaticArenaHeap() : ArenaHeap(_buf, heapsize) {}

 private:
  static const uint64_t heapsize = size - sizeof(ArenaHeap);
  uint8_t _buf[heapsize];
};

template <uint64_t size>
class StaticArena : public Arena {
 public:
  StaticArena() : Arena(&_heaps[0], &_heaps[1]) {}

 private:
  StaticArenaHeap<(size - sizeof(Arena)) / 2> _heaps[2];
};

template <class T>
  requires std::derived_from<T, PodObject>
Result<GcRoot<T>> GcRoot<T>::create(T* ptr, Arena& arena) {
  auto lst = TRY(arena.alloc<GcRootList>());
  arena.add_root(lst);
  return std::move(GcRoot<T>(ptr, lst));
}

template <class T>
  requires std::derived_from<T, PodObject>
GcRoot<T>::GcRoot(GcRoot<T>&& rhs) {
  rhs._node->update(this);
  _ptr = rhs._ptr;
  _node = rhs._node;
  rhs._ptr = 0;
  rhs._node = 0;
}

template <class T>
  requires std::derived_from<T, PodObject>
GcRoot<T>& GcRoot<T>::operator=(GcRoot<T>&& rhs) {
  if (_node != 0) _node->remove();

  rhs._node->update(this);
  _ptr = rhs._ptr;
  _node = rhs._node;
  rhs._ptr = 0;
  rhs._node = 0;

  return *this;
}