#pragma once

#include "arena.h"
/* ============================================================
   Helper functions
   ============================================================ */

/* Page size */
static u32 g_pagesize = 0;

static inline u32 arena_pagesize(void) {
  if (g_pagesize == 0) {
    g_pagesize = plat_get_pagesize();
  }
  return g_pagesize;
}

mem_arena *arena_block_from_pos(mem_arena *arena, u64 global_pos) {
  mem_arena *selected = arena;
  while (selected) {
    if (global_pos >= selected->base_pos && global_pos <= selected->pos) {
      return selected;
    }
    selected = selected->prev;
  }
  return NULL;
}

mem_arena *arena_block_from_index(mem_arena *arena, u64 index) {
  if (arena->push_size <= 0) {
    return NULL;
  }
  mem_arena *selected = arena;
  u64 pos = (index - 1) * arena->push_size + arena->block_index - 1;
  while (selected) {
    if (pos >= selected->base_pos && pos < selected->pos) {
      return selected;
    }
    selected = selected->prev;
  }
  return NULL;
}

mem_arena *arena_block_from_ptr(mem_arena *arena, u8 *ptr) {
  mem_arena *selected = arena;

  while (selected) {
    u8 *block_base =
        (u8 *)selected + ALIGN_UP_POW2(sizeof(mem_arena), selected->align);
    u8 *block_end = block_base + selected->pos - selected->base_pos;

    if (ptr >= block_base && ptr < block_end) {
      return selected;
    }
    selected = selected->prev;
  }
  return NULL;
}

u64 arena_pos_from_ptr(mem_arena *arena, void *ptr) {
  ASSERT(arena);
  ASSERT(ptr);
  if (!arena || !ptr) {
    return -1;
  }

  mem_arena *owner = arena_block_from_ptr(arena, ptr);
  ASSERT(owner);
  if (!owner) {
    return -1;
  }

  u8 *block_base = (u8 *)owner + ALIGN_UP_POW2(sizeof(mem_arena), owner->align);
  u64 local_offset = (u8 *)ptr - block_base;

  u64 pos = owner->base_pos + local_offset;

  if (pos >= owner->pos) {
    return -1;
  }
  return pos;
}

void *arena_ptr_from_pos(mem_arena *arena, u64 global_pos) {
  ASSERT(arena);
  ASSERT(global_pos >= 0);
  if (!arena || global_pos < 0) {
    return NULL;
  }

  mem_arena *owner = arena_block_from_pos(arena, global_pos);
  if (!owner) {
    return NULL;
  }

  u8 *block_base = (u8 *)owner + ALIGN_UP_POW2(sizeof(mem_arena), owner->align);
  u64 local_offset = global_pos - owner->base_pos;

  return (void *)(block_base + local_offset);
}

void *arena_ptr_from_index(mem_arena *arena, u64 index) {
  ASSERT(arena);
  ASSERT(index);
  if (!arena || !index || arena->push_size <= 0) {
    return 0;
  }

  mem_arena *owner = arena_block_from_index(arena, index);
  if (!owner) {
    return NULL;
  }

  u8 *block_base = (u8 *)owner + ALIGN_UP_POW2(sizeof(mem_arena), owner->align);
  u64 pos = (index - 1) * owner->push_size + owner->block_index - 1;
  u64 local_offset = pos - owner->base_pos;

  return (void *)(block_base + local_offset);
}

/* ============================================================
   Arena creation / destruction
   ============================================================ */

mem_arena *arena_create(arena_params *params) { // mk create
  ASSERT(params);
  ASSERT(params->reserve_size > 0);
  if (!params) {
    return NULL;
  }

  u32 pagesize = arena_pagesize();

  u64 reserve_size = ALIGN_UP_POW2(params->reserve_size, pagesize);
  u64 commit_size =
      params->commit_size ? ALIGN_UP_POW2(params->commit_size, pagesize) : 0;

  ASSERT(commit_size <= reserve_size);

  mem_arena *arena = (mem_arena *)plat_mem_reserve(reserve_size);
  if (!arena)
    return NULL;

  /* ------------------------------------------------------------
   Commit policy
   ------------------------------------------------------------ */

  u64 header_commit = ALIGN_UP_POW2(sizeof(mem_arena), pagesize);

  if (params->commit_policy == ARENA_COMMIT_FULL) {
    /* Commit everything */
    if (!plat_mem_commit(arena, reserve_size)) {
      plat_mem_release(arena, reserve_size);
      return NULL;
    }
    commit_size = reserve_size;
  } else {
    /* Lazy / partial commit */
    if (commit_size == 0 || commit_size < header_commit) {
      commit_size = header_commit;
    }

    if (!plat_mem_commit(arena, commit_size)) {
      plat_mem_release(arena, reserve_size);
      return NULL;
    }
  }

  /* Initialize header */
  arena->prev = NULL;
  arena->next = NULL;

  arena->base_pos = 0;
  arena->pos = 0;

  arena->reserve_size = reserve_size;
  arena->commit_size = commit_size;
  arena->commit_pos = commit_size;

  arena->align = params->align;
  arena->push_size = 0;

  arena->allow_free_list = params->allow_free_list;
  arena->free_list = NULL;
  arena->allow_swapback = 0;

  arena->growth_policy = params->growth_policy;
  arena->commit_policy = params->commit_policy;

  arena->max_nbre_blocks = params->max_nbre_blocks;
  arena->block_index = 1;

  if (arena->allow_free_list) {
    arena->push_size = ALIGN_UP_POW2(params->push_size, arena->align);

    if (!params->free_list) {
      arena->free_list = arena_create(&(arena_params){
          .reserve_size = MiB(1),
          .commit_size = MiB(1),
          .align = ARENA_ALIGN,
          .push_size = sizeof(arena_free_node),
          .allow_free_list = false,
          .free_list = NULL,
          .allow_swapback = true,
          .growth_policy = ARENA_GROWTH_NORMAL,
          .commit_policy = ARENA_COMMIT_LAZY,
          .max_nbre_blocks = 0,
      });

      if (!arena->free_list) {
        plat_mem_release(arena, arena->reserve_size);
        return NULL;
      }
    } else {
      arena->free_list = params->free_list;
    }
  } else {
    ASSERT(params->push_size > 0);
    arena->push_size = ALIGN_UP_POW2(params->push_size, arena->align);
    arena->allow_swapback = params->allow_swapback;
  }

  return arena;
}

void *arena_destroy(mem_arena **arena_ptr) { // mk destroy
  mem_arena *arena = *arena_ptr;
  if (!arena)
    return NULL;

  if (arena->free_list) {
    arena_destroy(&arena->free_list);
  }

  mem_arena *selected = arena;
  while (selected) {
    plat_mem_release(selected, selected->reserve_size);
    selected = selected->prev;
  }
  *arena_ptr = NULL;
  return (void *)1;
}

/* ============================================================
   Allocation
   ============================================================ */

void *arena_push(mem_arena **arena_ptr, u64 size, bool zero) { // mk push
  ASSERT(arena_ptr);
  if (!arena_ptr)
    return NULL;

  mem_arena *arena = *arena_ptr;

  /* ------------------------------------------------------------
     Resolve allocation size
     ------------------------------------------------------------ */

  if (arena->push_size > 0) {
    size = arena->push_size;
  } else {
    if (size == 0)
      return NULL;
    size = ALIGN_UP_POW2(size, arena->align);
  }

  /* ------------------------------------------------------------
     Free-list reuse (pointer mode)
     ------------------------------------------------------------ */

  if (arena->free_list) {
    mem_arena *fl = arena->free_list;
    u64 fl_base = ALIGN_UP_POW2(sizeof(mem_arena), fl->align);
    u64 count = fl->pos / fl->push_size;

    b32 arena_full = arena->max_nbre_blocks > 0 &&
                     arena->block_index >= arena->max_nbre_blocks;

    /* -------- Fixed-size: O(1) reuse -------- */
    if (arena->push_size > 0 && count > 0) {
      arena_free_node *node =
          (arena_free_node *)((u8 *)fl + fl_base +
                              (count - 1) * sizeof(arena_free_node));

      mem_arena *owner = arena_block_from_pos(arena, node->offset);
      ASSERT(owner);

      void *result = (u8 *)owner +
                     ALIGN_UP_POW2(sizeof(mem_arena), owner->align) +
                     (node->offset - owner->base_pos);

      arena_pop(&arena->free_list);

      if (zero)
        memset(result, 0, size);
      return result;
    }

    /* -------- Variable-size reuse -------- */
    if (arena->push_size == 0 && count > 0 &&
        (arena_full || arena->max_nbre_blocks == 0)) {

      for (mem_arena *selected_fl = fl; selected_fl;
           selected_fl = selected_fl->prev) {

        u64 selected_fl_count =
            (selected_fl->pos - selected_fl->base_pos) / selected_fl->push_size;

        for (u64 j = selected_fl_count; j >= 1; j--) {
          arena_free_node *node =
              (arena_free_node *)((u8 *)selected_fl + fl_base +
                                  (j - 1) * sizeof(arena_free_node));

          if (node->size >= size) {
            mem_arena *owner = arena_block_from_pos(arena, node->offset);
            ASSERT(owner);

            u8 *result = (u8 *)owner +
                         ALIGN_UP_POW2(sizeof(mem_arena), owner->align) +
                         (node->offset - owner->base_pos);

            u64 remaining = node->size - size;
            u64 rem_offset = node->offset + size;

            arena_swapback_pop(&arena->free_list, count);

            if (remaining > 0) {
              arena_free_node *rem =
                  (arena_free_node *)arena_push(&arena->free_list, 0, true);
              if (!rem)
                return NULL;

              assert(rem);

              rem->offset = rem_offset;
              rem->size = remaining;
            }

            if (zero)
              memset(result, 0, size);
            return result;
          }
          count--;
        }
      }
    }
  }

  /* ------------------------------------------------------------
     Normal allocation (last block)
     ------------------------------------------------------------ */

  mem_arena *selected = arena;

  u64 local_pos = selected->pos - selected->base_pos;
  u64 local_pre = ALIGN_UP_POW2(local_pos, selected->align);
  u64 local_post = local_pre + size;

  if (local_post > selected->reserve_size -
                       ALIGN_UP_POW2(sizeof(mem_arena), selected->align)) {

    if (arena->allow_free_list && arena->push_size == 0) {
      u64 tail_start = selected->pos;
      u64 tail_size = selected->reserve_size -
                      (tail_start - selected->base_pos +
                       ALIGN_UP_POW2(sizeof(mem_arena), selected->align));

      if (tail_size > 0) {
        arena_free_node *node =
            (arena_free_node *)arena_push(&arena->free_list, 0, true);
        assert(node);
        if (!node)
          return NULL;

        node->offset = tail_start;
        node->size = tail_size;
        selected->pos = selected->base_pos + selected->reserve_size;
      }
    }

    if (selected->next) {
      selected->next->base_pos = selected->pos + 1;
      selected->next->pos = selected->pos + 1;
      selected->next->free_list = selected->free_list;
      selected = selected->next;
      *arena_ptr = selected;
      local_pre = 0;
      local_post = size;
    } else {

      /* ------------------------------------------------------------
         Grow arena if needed
         ------------------------------------------------------------ */

      if (arena->max_nbre_blocks &&
          arena->block_index >= arena->max_nbre_blocks) {
        printf("Arena full.\n");
        return NULL;
      }

      u64 new_reserve = selected->reserve_size;
      if (arena->growth_policy == ARENA_GROWTH_DOUBLE)
        new_reserve *= 2;

      arena_params p = {
          .reserve_size = new_reserve,
          .commit_size = selected->commit_size,
          .align = selected->align,
          .push_size = arena->push_size,
          .allow_free_list = arena->allow_free_list,
          .free_list = arena->free_list,
          .allow_swapback = arena->allow_swapback,
          .growth_policy = arena->growth_policy,
          .commit_policy = arena->commit_policy,
          .max_nbre_blocks = arena->max_nbre_blocks,
      };

      mem_arena *next = arena_create(&p);
      if (!next)
        return NULL;

      next->base_pos = selected->pos + 1;
      next->pos = selected->pos + 1;
      next->prev = selected;
      selected->next = next;
      next->block_index = selected->block_index + 1;

      selected = next;
      *arena_ptr = selected;

      local_pre = 0;
      local_post = size;
    }
  }

  /* ------------------------------------------------------------
     Commit memory if needed
     ------------------------------------------------------------ */

  if (local_post > selected->commit_pos) {
    u64 new_commit = ALIGN_UP_POW2(local_post, arena_pagesize());
    new_commit = MIN(new_commit, selected->reserve_size);

    if (!plat_mem_commit((u8 *)selected + selected->commit_pos,
                         new_commit - selected->commit_pos)) {
      return NULL;
    }

    selected->commit_pos = new_commit;
  }

  /* ------------------------------------------------------------
     Finalize allocation
     ------------------------------------------------------------ */

  u8 *result = (u8 *)selected +
               ALIGN_UP_POW2(sizeof(mem_arena), selected->align) + local_pre;

  selected->pos = selected->base_pos + local_post;

  if (zero)
    memset(result, 0, size);

  return result;
}

/* ============================================================
   Free (pointer mode): pop to free list
   ============================================================ */

void *arena_free(mem_arena **arena_ptr, u8 **ptr, u64 size) { // mk free

  mem_arena *arena = *arena_ptr;

  ASSERT(arena);
  ASSERT(arena->allow_free_list);
  ASSERT(ptr && *ptr);

  if (!arena || !arena->allow_free_list || !ptr || !*ptr)
    return NULL;

  u64 elem_size = arena->push_size ? arena->push_size : size;
  ASSERT(elem_size > 0);

  /* ------------------------------------------------------------
     Find owning block
     ------------------------------------------------------------ */

  mem_arena *selected = arena;
  mem_arena *owner = arena_block_from_ptr(arena, *ptr);
  ASSERT(owner);
  if (!owner) {
    return NULL;
  }

  /* ------------------------------------------------------------
     Compute global offset using arena_pos()
     ------------------------------------------------------------ */

  u64 global_offset = arena_pos_from_ptr(arena, *ptr);
  if (global_offset == -1) {
    return NULL;
  }

  /* ------------------------------------------------------------
     Fast path: pop only if this is the LAST block
     ------------------------------------------------------------ */

  if (owner == arena && global_offset + elem_size == arena->pos) {
    arena->pos -= elem_size;
    if (arena->pos < arena->base_pos) {
      arena->prev->pos = arena->pos - 1;
      arena->pos = arena->base_pos;
      arena->prev->free_list = arena->free_list;
      *arena_ptr = arena->prev;
    }
    *ptr = NULL;
    return (void *)1;
  }

  /* ------------------------------------------------------------
     Otherwise push into free list
     ------------------------------------------------------------ */

  arena_free_node *node = (arena_free_node *)arena_push(
      &arena->free_list, sizeof(arena_free_node), false);
  if (!node)
    return NULL;

  node->offset = global_offset;
  node->size = elem_size;

  *ptr = NULL;
  return (void *)1;
}

/* ============================================================
   Stack operations
   ============================================================ */

void *arena_pop_to(mem_arena **arena_ptr, u64 count) { // mk pop to
  ASSERT(arena_ptr);

  mem_arena *arena = *arena_ptr;

  ASSERT(!arena->allow_free_list);
  ASSERT(arena->push_size > 0);
  if (arena->allow_free_list || arena->push_size <= 0) {
    return NULL;
  }

  u64 target_pos;
  if (arena->pos < count * arena->push_size) {
    target_pos = 0;
  } else {
    target_pos = arena->pos - count * arena->push_size;
  }

  mem_arena *selected = arena;

  while (true) {
    if (selected->base_pos <= target_pos) {
      selected->pos = target_pos;
      selected->free_list = arena->free_list;
      break;
    }
    selected->pos = selected->base_pos;
    target_pos = target_pos - 1;
    selected = selected->prev;
  }
  *arena_ptr = selected;

  return (void *)1;
}

void *arena_swapback_pop(mem_arena **arena_ptr, u64 index) { // mk swapback

  mem_arena *arena = *arena_ptr;

  ASSERT(arena);
  ASSERT(arena->push_size > 0);
  ASSERT(arena->allow_swapback);
  if (arena->push_size <= 0 || !arena->allow_swapback) {
    return NULL;
  }

  u64 count = arena->pos / arena->push_size;
  ASSERT(index <= count);
  if (index > count) {
    return NULL;
  }

  /* Last element: just pop */
  if (index == count - 1) {
    u8 *r = arena_pop(&arena);

    *arena_ptr = arena;

    return r;
  }

  mem_arena *owner = arena_block_from_index(arena, index);

  if (!owner) {
    fprintf(stderr, "ERROR: Swapback pop failed, index out of range");
    return NULL;
  }
  u8 *owner_base = (u8 *)owner + ALIGN_UP_POW2(sizeof(mem_arena), owner->align);
  u8 *arena_base = (u8 *)arena + ALIGN_UP_POW2(sizeof(mem_arena), arena->align);

  u8 *dst = arena_ptr_from_index(arena, index);
  u8 *src = arena_ptr_from_index(arena, count);

  memcpy(dst, src, arena->push_size);

  u8 *r = arena_pop(&arena);

  *arena_ptr = arena;

  return r;
}

/* ============================================================
   Utilities
   ============================================================ */

void *arena_clear(mem_arena **arena_ptr) { // mk clear

  mem_arena *arena = *arena_ptr;

  if (!arena)
    return NULL;

  mem_arena *selected = arena;
  while (selected->prev) {
    selected->pos = 0;
    selected->base_pos = 0;
    selected = selected->prev;
  }
  selected->pos = 0;
  selected->free_list = arena->free_list;

  *arena_ptr = selected;

  if (arena->free_list) {
    if (!arena_clear(&arena->free_list))
      return NULL;
  }
  return (void *)1;
}

mem_arena *arena_merge(mem_arena **dst_ptr, mem_arena **src_ptr) { // mk merge

  mem_arena *dst = *dst_ptr;
  mem_arena *src = *src_ptr;

  if (!dst || !src || dst == src)
    return NULL;

  /* ------------------------------------------------------------
     Config compatibility
     ------------------------------------------------------------ */

  if (dst->align != src->align || dst->push_size != src->push_size ||
      dst->allow_free_list != src->allow_free_list ||
      dst->allow_swapback != src->allow_swapback ||
      dst->growth_policy != src->growth_policy) {
    return NULL;
  }

  /* ------------------------------------------------------------
     Merge free lists
     ------------------------------------------------------------ */
  u64 block_base_pos = dst->pos + 1;

  if (dst->allow_free_list) {
    u64 fl_base = ALIGN_UP_POW2(sizeof(mem_arena), src->free_list->align);

    for (mem_arena *selected_fl = src->free_list; selected_fl;
         selected_fl = selected_fl->prev) {

      u64 selected_fl_count =
          (selected_fl->pos - selected_fl->base_pos) / selected_fl->push_size;

      for (u64 j = selected_fl_count; j >= 1; j--) {
        arena_free_node *node =
            (arena_free_node *)((u8 *)selected_fl + fl_base +
                                (j - 1) * sizeof(arena_free_node));
        node->offset += block_base_pos;
      }
    }

    src->free_list = arena_merge(&dst->free_list, &src->free_list);
    if (!src->free_list)
      return NULL;
  }

  /* ------------------------------------------------------------
     Walk src blocks once:
     ------------------------------------------------------------ */

  mem_arena *src_first = NULL;

  for (mem_arena *b = src; b; b = b->prev) {
    src_first = b;
  }
  /* ------------------------------------------------------------
     Update global metadata
     ------------------------------------------------------------ */

  u8 selected_block_index = dst->block_index;
  dst->max_nbre_blocks += src->max_nbre_blocks;

  mem_arena *selected = src_first;
  while (selected) {
    selected->pos += block_base_pos;
    selected->base_pos += block_base_pos;
    selected->block_index += selected_block_index;
    selected->max_nbre_blocks = dst->max_nbre_blocks;
    selected = selected->next;
  }

  for (mem_arena *a = dst; a; a = a->prev) {
    a->max_nbre_blocks = dst->max_nbre_blocks;
  }

  if (dst->next) {
    mem_arena *src_last = NULL;
    for (mem_arena *b = src; b; b = b->next) {
      src_last = b;
    }

    mem_arena *dst_fst_empty_block = dst->next;

    mem_arena *dst_empty_block = dst_fst_empty_block;
    u8 dst_empty_block_index = src_last->block_index;
    while (dst_empty_block) {
      dst_empty_block->block_index = dst_empty_block_index + 1;
      dst_empty_block->max_nbre_blocks = dst->max_nbre_blocks;
      dst_empty_block = dst_empty_block->next;
    }

    /* ------------------------------------------------------------
       Stitch block chains
       ------------------------------------------------------------ */

    src_last->next = dst_fst_empty_block;
    dst_fst_empty_block->prev = src_last;
  }

  dst->next = src_first;
  src_first->prev = dst;

  *dst_ptr = src;
  *src_ptr = NULL;

  return *dst_ptr;
}

/* ============================================================
   Temp arenas
   ============================================================ */

mem_arena_temp arena_temp_begin(mem_arena *arena) {
  ASSERT(arena);
  ASSERT(!arena->allow_free_list);
  return (mem_arena_temp){arena, arena->pos};
}

void arena_temp_end(mem_arena_temp temp) {
  ASSERT(temp.arena);
  ASSERT(!temp.arena->allow_free_list);
  arena_pop_to(&temp.arena, temp.pos / temp.arena->push_size);
}

static THREAD_LOCAL mem_arena *_scratch_arenas[2] = {NULL, NULL};

mem_arena_temp arena_scratch_get(mem_arena **conflicts, u32 num_conflicts) {
  i32 scratch_index = -1;

  for (i32 i = 0; i < 2; i++) {
    b32 conflict_found = false;

    for (u32 j = 0; j < num_conflicts; j++) {
      if (_scratch_arenas[i] == conflicts[j]) {
        conflict_found = true;
        break;
      }
    }

    if (!conflict_found) {
      scratch_index = i;
      break;
    }
  }

  if (scratch_index == -1) {
    return (mem_arena_temp){0};
  }

  mem_arena **selected = &_scratch_arenas[scratch_index];

  if (*selected == NULL) {
    arena_params params = {
        .reserve_size = MiB(64),
        .commit_size = MiB(1),
        .align = ARENA_ALIGN,
        .push_size = 8,
        .allow_free_list = false,
        .allow_swapback = true,
        .growth_policy = ARENA_GROWTH_NORMAL,
        .commit_policy = ARENA_COMMIT_LAZY,
        .max_nbre_blocks = 0,
    };

    *selected = arena_create(&params);
  }

  return arena_temp_begin(*selected);
}

void arena_scratch_release(mem_arena_temp scratch) { arena_temp_end(scratch); }