#include "allocator.h"

static MemoryAllocator* mem_alloc = nullptr;

MemoryAllocator::MemoryAllocator(uint32_t start, size_t size) {
  if (size < sizeof(MemoryChunk)) {
    begin = nullptr;
    return;
  }
  begin = new (reinterpret_cast<void*>(start)) MemoryChunk{.size = static_cast<size_t>(size - sizeof(MemoryChunk))};
  if (mem_alloc == nullptr) mem_alloc = this;
}

void* MemoryAllocator::malloc(size_t size) {
  // TODO spinlock
  // TODO memset 0 over allocated memory
  MemoryChunk* result = nullptr;

  // find the first chunk that is at least size and not allocated
  for (auto* iter = begin; iter != nullptr; iter = iter->next) {
    if (iter->size >= size && !iter->allocated) {
      result = iter;

      // split chunk if there is at least one byte left over
      if (result->size >= size + sizeof(MemoryChunk) + 1) {
        auto* _new = new (result + sizeof(MemoryChunk) + size) MemoryChunk{
            .prev = result, .next = result->next, .size = static_cast<size_t>(result->size - sizeof(MemoryChunk))};
        result->next = _new;
      }

      break;
    }
  }

  return result + sizeof(MemoryChunk);
}
void MemoryAllocator::free(void* ptr) {
  auto* chunk = reinterpret_cast<MemoryChunk*>(reinterpret_cast<uint32_t>(ptr) - sizeof(MemoryChunk));
  chunk->allocated = false;

  // merge in next chunk
  if (chunk->next != nullptr && !chunk->next->allocated) {
    chunk->size += chunk->next->size + sizeof(MemoryChunk);
    chunk->next = chunk->next->next;
    if (chunk->next != nullptr) chunk->next->prev = chunk;
  }

  // merge into previous chunk
  if (chunk->prev != nullptr && !chunk->prev->allocated) {
    chunk->prev->size += chunk->size + sizeof(MemoryChunk);
    chunk->prev->next = chunk->next;
    if (chunk->next != nullptr) chunk->next->prev = chunk->prev;
  }
};

size_t MemoryAllocator::free_space() const {
  size_t r = 0;
  for (auto* iter = begin; iter != nullptr; iter = iter->next)
    if (!iter->allocated) r += iter->size;
  return r;
}
size_t MemoryAllocator::used_space() const {
  size_t r = 0;
  for (auto* iter = begin; iter != nullptr; iter = iter->next)
    if (iter->allocated) r += iter->size;
  return r;
}

void* operator new(unsigned sz, const std::nothrow_t&) noexcept {
  if (mem_alloc == nullptr) return nullptr;
  return mem_alloc->malloc(sz);
}
void* operator new[](unsigned sz, const std::nothrow_t&) noexcept {
  if (mem_alloc == nullptr) return nullptr;
  return mem_alloc->malloc(sz);
}

void operator delete(void* ptr) {
  if (mem_alloc != nullptr) mem_alloc->free(ptr);
}
void operator delete[](void* ptr) {
  if (mem_alloc != nullptr) mem_alloc->free(ptr);
}

void* operator new(unsigned /* size_bytes */, void* ptr) {
  return ptr;
}
void* operator new[](unsigned /* size_bytes */, void* ptr) {
  return ptr;
}