#include "mem_arena.h"
#include "aliases.h"
#include "mem_allocator.h"
#include "mem_utils.h"
#include <math.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

#ifndef DEFAULT_ALIGNMENT
// Why 2 * sizeof(void *) instead of sizeof(void *)
// https://handmade.network/forums/t/6860-alignment_arena_allocator
#define DEFAULT_ALIGNMENT (2 * sizeof(void *))
#endif /* ifndef DEFAULT_ALIGNMENT */

#define HDR_MAGIC_BYTE_COUNT 8
#define HDR_MAGIC                                                              \
  { 0x57, 0x41, 0x41, 0x52, 0x4e, 0x48, 0x44, 0x52 }
#define MAX_HDR_SEARCH_LENGTH 256

typedef struct arena_alloc_hdr ArenaAllocHDR;
struct arena_alloc_hdr {
  u8 magic[HDR_MAGIC_BYTE_COUNT];
  u64 alloc_size;
  u64 alignment;
  u8 *alloc_start;
};

typedef struct base_arena BaseArena;
struct base_arena {
  u8 *buf;
  u8 *offset;
  u64 capacity;
  BaseArena *prev;
  BaseArena *next;
};

struct growing_arena {
  BaseArena *active_arena;
  u64 count;
  u64 initial_capacity;
};

internal BaseArena *find_arena_from_pointer(const Arena *arena, void *ptr);

internal bool base_arena_init(BaseArena *arena, u64 capacity);
internal void *base_arena_alloc_aligned(BaseArena *arena, u64 size,
                                        u64 alignment);
internal void base_arena_clear(BaseArena *arena);
internal void base_arena_free(BaseArena *arena);
internal ArenaAllocHDR *find_alloc_header(BaseArena *arena, void *alloc_ptr);

internal void *mem_arena_alloc(u64 size, void *alloc_obj);
internal void *mem_arena_alloc_aligned(u64 size, u64 alignment,
                                       void *alloc_obj);
internal void *mem_arena_realloc(void *ptr, u64 size, void *alloc_obj);
internal void *mem_arena_realloc_aligned(void *ptr, u64 size, u64 alignment,
                                         void *alloc_obj);

// PUBLIC API

Allocator wapp_mem_arena_allocator(const Arena *arena) {
  return (Allocator){
      .obj = (void *)arena,
      .alloc = mem_arena_alloc,
      .alloc_aligned = mem_arena_alloc_aligned,
      .realloc = mem_arena_realloc,
      .realloc_aligned = mem_arena_realloc_aligned,
      .free = NULL,
  };
}

bool wapp_mem_arena_init(Arena **arena, u64 base_capacity) {
  if (!arena || *arena || base_capacity == 0) {
    return false;
  }

  *arena = (Arena *)calloc(1, sizeof(Arena));
  Arena *arena_ptr = *arena;
  if (!arena_ptr) {
    return false;
  }

  arena_ptr->active_arena = (BaseArena *)calloc(1, sizeof(BaseArena));
  if (!(arena_ptr->active_arena)) {
    wapp_mem_arena_free(arena);
    return false;
  }

  if (!base_arena_init(arena_ptr->active_arena, base_capacity)) {
    wapp_mem_arena_free(arena);
    return false;
  }

  arena_ptr->count = 1;
  arena_ptr->initial_capacity = base_capacity;

  return true;
}

void *wapp_mem_arena_alloc(Arena *arena, u64 size) {
  return wapp_mem_arena_alloc_aligned(arena, size, DEFAULT_ALIGNMENT);
}

void *wapp_mem_arena_alloc_aligned(Arena *arena, u64 size, u64 alignment) {
  if (!arena || !(arena->active_arena)) {
    return NULL;
  }

  void *output = base_arena_alloc_aligned(arena->active_arena, size, alignment);
  if (!output) {
    if (arena->active_arena->next) {
      arena->active_arena = arena->active_arena->next;
    } else {
      arena->active_arena->next = (BaseArena *)calloc(1, sizeof(BaseArena));
      if (!(arena->active_arena->next)) {
        return NULL;
      }

      if (!base_arena_init(arena->active_arena->next,
                           arena->initial_capacity)) {
        free(arena->active_arena->next);
        return NULL;
      }

      arena->active_arena->next->prev = arena->active_arena;
      arena->active_arena = arena->active_arena->next;

      ++(arena->count);
    }

    output = base_arena_alloc_aligned(arena->active_arena, size, alignment);
    if (!output) {
      return NULL;
    }
  }

  memset(output, 0, size);

  return output;
}

void *wapp_mem_arena_realloc(Arena *arena, void *ptr, u64 size) {
  return wapp_mem_arena_realloc_aligned(arena, ptr, size, DEFAULT_ALIGNMENT);
}

void *wapp_mem_arena_realloc_aligned(Arena *arena, void *ptr, u64 size,
                                     u64 alignment) {
  if (!arena) {
    return NULL;
  }

  BaseArena *base_arena = find_arena_from_pointer(arena, ptr);
  if (!base_arena) {
    return NULL;
  }

  ArenaAllocHDR *header = find_alloc_header(base_arena, ptr);
  if (!header) {
    return NULL;
  }

  if (header->alloc_start + header->alloc_size == base_arena->offset) {
    // Avoid allocating new pointer and copying memory if pointer is at the end
    // of the arena
    i64 diff = size - header->alloc_size;

    u8 *new_offset = base_arena->offset + diff;

    u8 *clear_start = diff < 0 ? new_offset : base_arena->offset;
    memset(clear_start, 0, llabs(diff));

    header->alloc_size = size;
    base_arena->offset = new_offset;

    return header->alloc_start;
  }

  void *new_alloc = wapp_mem_arena_alloc_aligned(arena, size, alignment);
  if (!new_alloc) {
    return NULL;
  }

  u64 to_copy = size < header->alloc_size ? size : header->alloc_size;

  memcpy(new_alloc, ptr, to_copy);

  return new_alloc;
}

void wapp_mem_arena_clear(Arena *arena) {
  if (!arena) {
    return;
  }

  BaseArena *last_active = NULL;
  while (arena->active_arena) {
    base_arena_clear(arena->active_arena);

    last_active = arena->active_arena;

    arena->active_arena = arena->active_arena->prev;
  }

  arena->active_arena = last_active;
}

void wapp_mem_arena_free(Arena **arena) {
  if (!arena) {
    return;
  }

  Arena *arena_ptr = *arena;
  if (!arena_ptr) {
    return;
  }

  BaseArena *current;
  BaseArena *next;
  BaseArena *prev;

  current = arena_ptr->active_arena->next;
  while (current) {
    next = current->next;

    base_arena_free(current);
    free(current);

    current = next;
  }

  current = arena_ptr->active_arena->prev;
  while (current) {
    prev = current->prev;

    base_arena_free(current);
    free(current);

    current = prev;
  }

  base_arena_free(arena_ptr->active_arena);

  free(arena_ptr->active_arena);
  arena_ptr->active_arena = NULL;

  arena_ptr->count = 0;
  arena_ptr->initial_capacity = 0;

  free(*arena);
  *arena = NULL;
}

// INTERNAL FUNCTIONS

internal BaseArena *find_arena_from_pointer(const Arena *arena, void *ptr) {
  if (!arena || !ptr) {
    return NULL;
  }

  // Ensure pointer is not out of bounds
  u8 *alloc = (u8 *)ptr;
  BaseArena *active = arena->active_arena;

  if (alloc > active->buf + arena->initial_capacity) {
    return NULL;
  }

  for (u64 i = 0; i < arena->count; ++i) {
    if (alloc >= active->buf && alloc < active->buf + arena->initial_capacity) {
      return (BaseArena *)active;
    }

    active = active->prev;
  }

  return NULL;
}

internal bool base_arena_init(BaseArena *arena, u64 capacity) {
  if (!arena || arena->buf || capacity == 0) {
    return false;
  }

  arena->buf = (u8 *)calloc(capacity, sizeof(u8));
  if (!(arena->buf)) {
    return false;
  }

  arena->capacity = capacity;
  arena->offset = arena->buf;
  arena->prev = arena->next = NULL;

  return true;
}

internal void *base_arena_alloc_aligned(BaseArena *arena, u64 size,
                                        u64 alignment) {
  if (!arena) {
    return NULL;
  }

  u8 *start_offset = arena->offset;
  u8 *alloc_start = arena->offset + sizeof(ArenaAllocHDR);

  u8 *output = wapp_mem_util_align_forward((void *)alloc_start, alignment);
  if (output + size >= arena->buf + arena->capacity) {
    return NULL;
  }

  ArenaAllocHDR *header = (ArenaAllocHDR *)start_offset;
  *header = (ArenaAllocHDR){
      .magic = HDR_MAGIC,
      .alloc_size = size,
      .alignment = alignment,
      .alloc_start = output,
  };

  arena->offset = output + size;

  return (void *)output;
}

internal void base_arena_clear(BaseArena *arena) {
  if (!arena) {
    return;
  }

  memset(arena->buf, 0, arena->offset - arena->buf);
  arena->offset = arena->buf;
}

internal void base_arena_free(BaseArena *arena) {
  if (!arena) {
    return;
  }

  if (arena->buf) {
    free(arena->buf);
  }

  arena->buf = arena->offset = NULL;
  arena->capacity = 0;
  arena->prev = arena->next = NULL;
}

internal ArenaAllocHDR *find_alloc_header(BaseArena *arena, void *alloc_ptr) {
  persistent const u8 magic[HDR_MAGIC_BYTE_COUNT] = HDR_MAGIC;

  u8 *current = (u8 *)alloc_ptr;

  u8 *max_search_end = current - MAX_HDR_SEARCH_LENGTH;
  u8 *arena_buf_start = arena->buf;

  u8 *search_end =
      max_search_end > arena_buf_start ? max_search_end : arena_buf_start;

  bool match;
  for (; current >= search_end; --current) {
    match = true;

    for (u64 i = 0; i < HDR_MAGIC_BYTE_COUNT; ++i) {
      if (current[i] != magic[i]) {
        match = false;
        break;
      }
    }

    if (match) {
      return (ArenaAllocHDR *)current;
    }
  }

  return NULL;
}

internal void *mem_arena_alloc(u64 size, void *alloc_obj) {
  Arena *arena = (Arena *)alloc_obj;

  return wapp_mem_arena_alloc(arena, size);
}

internal void *mem_arena_alloc_aligned(u64 size, u64 alignment,
                                       void *alloc_obj) {
  Arena *arena = (Arena *)alloc_obj;

  return wapp_mem_arena_alloc_aligned(arena, size, alignment);
}

internal void *mem_arena_realloc(void *ptr, u64 size, void *alloc_obj) {
  Arena *arena = (Arena *)alloc_obj;

  return wapp_mem_arena_realloc(arena, ptr, size);
}

internal void *mem_arena_realloc_aligned(void *ptr, u64 size, u64 alignment,
                                         void *alloc_obj) {
  Arena *arena = (Arena *)alloc_obj;

  return wapp_mem_arena_realloc_aligned(arena, ptr, size, alignment);
}