/* Copyright (C) 2025 Aryadev Chavali

 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License Version 2 for
 * details.

 * You may distribute and modify this code under the terms of the GNU General
 * Public License Version 2, which you should have received a copy of along with
 * this program.  If not, please go to <https://www.gnu.org/licenses/>.

 * Created: 2025-04-05
 * Description: Implementations for memory models.
 */

#include <lib/arena.h>

#include <malloc.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

page_t *page_create(u64 size)
{
  size           = MAX(size, PAGE_DEFAULT_SIZE);
  page_t *page   = calloc(1, sizeof(*page) + size);
  page->next     = NULL;
  page->size     = 0;
  page->capacity = size;
  return page;
}

void page_resize(page_t **page, u64 size)
{
  if (!page)
    return;
  else if (!(*page))
    *page = page_create(size);
  else if (page[0]->capacity < size)
  {
    page[0]->capacity = MAX(size, page[0]->capacity * 1.5);
    page[0]           = realloc(page[0], sizeof(*page[0]) + page[0]->capacity);
  }
}

i64 page_append(page_t *page, void *data, u64 size)
{
  if (!page || page->size + size >= page->capacity)
    return -1;
  if (data)
    memcpy(page->data + page->size, data, size);
  u64 ptr = page->size;
  page->size += size;
  return ptr;
}

u64 page_rappend(page_t **page, void *data, u64 size)
{
  page_resize(page, page[0]->size + size);
  if (data)
    memcpy(page[0]->data + page[0]->size, data, size);
  u64 ptr = page[0]->size;
  page[0]->size += size;
  return ptr;
}

void *arena_alloc(arena_t *arena, u64 size)
{
  if (!arena)
    return NULL;
  else if (!arena->start)
  {
    arena->start = page_create(1);
    arena->end   = arena->start;
  }

  page_t *best_fit;
  for (best_fit = arena->start; best_fit; best_fit = best_fit->next)
    if (best_fit->size + size + MEMORY_ALIGNMENT < best_fit->capacity)
      break;

  if (!best_fit)
  {
    best_fit         = page_create(size);
    arena->end->next = best_fit;
    arena->end       = best_fit;
  }

  // NOTE: Fsanitize has a hissy fit if we don't align memory "correctly"
  u64 offset = 0;

  if (size % MEMORY_ALIGNMENT == 0)
  {
    u64 div = (((u64)(best_fit->data + best_fit->size)) % MEMORY_ALIGNMENT);
    if (div != 0)
      offset = MEMORY_ALIGNMENT - div;
  }

  void *start = best_fit->data + best_fit->size + offset;

  best_fit->size += size + offset;

  return start;
}

void *arena_realloc(arena_t *arena, void *ptr, u64 oldsize, u64 newsize)
{
  if (!ptr)
    return arena_alloc(arena, newsize);
  else if (!arena || !arena->start)
    return NULL;
  else if (newsize <= oldsize)
    // No need to change anything.
    return ptr;

  bool copy_into_new = true;
  void *start        = NULL;
  page_t *old_page = NULL, *best_fit = NULL;

  for (page_t *page = arena->start; page; page = page->next)
  {
    if (!best_fit && page->size + newsize < page->capacity)
      best_fit = page;
    if (page->data <= (u8 *)ptr &&
        (u8 *)(ptr) + oldsize <= page->data + page->capacity)
      old_page = page;
  }

  // If the old page exists, ptr is the latest allocation in it, and it has
  // enough space to contain the new size, just resize and return.
  if (old_page && old_page->data + old_page->size - oldsize == ptr &&
      old_page->size - oldsize + newsize < old_page->capacity)
  {
    start = ptr;
    old_page->size += newsize - oldsize;
    copy_into_new = false;
  }
  else
  {
    if (!old_page)
      copy_into_new = false;
    if (!best_fit)
    {
      best_fit         = page_create(newsize);
      arena->end->next = best_fit;
      arena->end       = best_fit;
    }

    start = best_fit->data + best_fit->size;
    best_fit->size += newsize;
  }

  if (copy_into_new)
  {
    memcpy(start, ptr, oldsize);
    memset(start + oldsize, 0, newsize - oldsize);
  }

  return start;
}

void arena_attach(arena_t *arena, page_t *page)
{
  if (!arena || !page)
    return;
  else if (!arena->start || !arena->end)
  {
    arena->start = page;
    arena->end   = page;
  }
  else
  {
    page->next   = arena->start;
    arena->start = page;
  }
}

void arena_reset(arena_t *arena)
{
  if (!arena || !arena->start)
    return;

  for (page_t *cur = arena->start; cur; cur = cur->next)
  {
    if (cur->size == 0)
      continue;
    cur->size = 0;
    memset(cur->data, 0, cur->capacity);
  }
}

void arena_cleanup(arena_t *arena)
{
  if (!arena || !arena->start)
    return;

  for (page_t *cur = arena->start, *next = NULL; cur;
       next = cur->next, free(cur), cur = next)
    continue;

  memset(arena, 0, sizeof(*arena));
}