#add_context arena      := Arena.{ proc = PagingArenaProc };
#add_context temp_arena := Arena.{ proc = TempArenaProc   };

/// Arena is an interface for an arena-based memory allocation.
Arena :: struct {
   proc: ArenaProc;
   data: *void;
}

ArenaProc :: #type (
   event:       ArenaEvent,
   arena_data: *void,

   new_size: int, old_size: int, align: int,
   new_ptr: *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void;

ArenaEvent :: enum {
   Setup;
   Teardown;

   Acquire;
   Reset;

   Save;
   Restore;

   Report;
}

/// ArenaPush pushes a value of type T to the current context.arena.
ArenaPush :: ($T: Type, loc := #caller_location) -> *T {
   return ArenaPushSize(size_of(T), align_of(T), loc = loc).(*T);
}

/// ArenaPushSize pushes size bytes to the current context.arena.
ArenaPushSize :: (size: int, align := DefaultAlign, loc := #caller_location) -> *void {
   if size < 0
    { Panic("jc: size cannot be negative", loc = loc); }
   TrySetupArena(*context.arena);
   return context.arena.proc(.Acquire, context.arena.data, size, 0, align, null, null, loc);
}

ArenaSave :: (loc := #caller_location) -> ArenaSaveResult {
   TrySetupArena(*context.arena, loc);
   return context.arena.proc(.Save, context.arena.data, 0, 0, 0, null, null, loc).(ArenaSaveResult);
}

ArenaRestore :: (savepoint: ArenaSaveResult, loc := #caller_location) {
   context.arena.proc(.Restore, context.arena.data, 0, 0, 0, null, savepoint.(*void), loc);
}

ArenaAutoRestore :: () #expand {
   wm := ArenaSave();
   `defer ArenaRestore(wm);
}

/// ArenaSaveResult is an arena-specific value used for saving/restoring.
/// see: ArenaSave
ArenaSaveResult :: #type,distinct *void;


/// ArenaReset resets arena's state, allowing its memory to be reused.
ArenaReset :: (loc := #caller_location) {
   TrySetupArena(*context.arena);
   context.arena.proc(.Reset, context.arena.data, 0, 0, 0, null, null, loc);
}

/// ArenaRelease frees the current context.arena's memory.
ArenaRelease :: (loc := #caller_location) {
   TrySetupArena(*context.arena, loc);
   context.arena.proc(.Teardown, context.arena.data, 0, 0, 0, null, null, loc);
}

ArenaReport :: (loc := #caller_location) {
   TrySetupArena(*context.arena, loc);
   context.arena.proc(.Report, context.arena.data, 0, 0, 0, null, null, loc);
}

/// ArenaToAllocator wraps the given arena, allowing it to be used with Jai's builtin Allocator system.
ArenaToAllocator :: (arena: *Arena) -> Allocator {
   return .{ proc = JaiAllocatorProc, data = arena };
}

/// PanicArena is an Arena that panics when used.
PanicArena :: () -> Arena {
   return .{ proc = PanicArenaProc };
}

BumpData :: struct {
   memory: []u8;
   offset: int;
}

/// BumpArena is a simple linear Arena.
BumpArena :: (bump: *BumpData) -> Arena {
   return .{ proc = BumpArenaProc, data = bump };
}

PagingArena :: () -> Arena {
   return .{ proc = PagingArenaProc };
}


#scope_module

TrySetupArena :: (arena: *Arena, loc := #caller_location) #expand {
   if arena.data != null
    { return; }

   data := arena.proc(.Setup, null, 0, 0, 0, null, null, loc);
   if data != null
    { arena.data = data; }
} @jc.nodocs

PanicArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   if event == {
      case .Acquire;
         if new_size > 0
          { Panic("jc: PanicArena, acquiring memory is not allowed", loc = caller); }

      case .Save;
         Panic("jc: PanicArena, saving is not allowed", loc = caller);
      case .Restore;
         Panic("jc: PanicArena, restoring is not allowed", loc = caller);
   }

   return null;
} @jc.nodocs

BumpArenaProc :: (
   event:       ArenaEvent,
   arena_data: *void,

   new_size: int, old_size: int, align: int,
   new_ptr: *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void #no_abc {
   bump := arena_data.(*BumpData);
   if event == {
      case .Setup;
         if bump.memory.data == null
          { Panic("jc: BumpArena, no backing memory", loc = caller); }
         return arena_data;

      case .Teardown;
         // Do nothing, we don't own this memory

      case .Acquire;
         end := bump.memory.data + bump.offset;
         if new_size == 0
          { return end; }

         ptr  := MemAlignForward(end, align);
         size := new_size + (ptr - end);
         if bump.offset + size > bump.memory.count
          { Panic("jc: BumpArena, out of memory", loc = caller); }

         bump.offset += size;
         return ptr;
      case .Reset;
         bump.offset = 0;

      case .Save;
         return bump.offset.(*void);
      case .Restore;
         wm := old_ptr.(int);
         if wm < 0 || wm >= bump.memory.count
          { Panic("jc: BumpArena, restored invalid savepoint", loc = caller); }
         bump.offset = wm;

      case .Report;
         WriteString("\n--- BumpArena Report ---");

         WriteString("\nMem Max: ");
         WriteNumber(bump.memory.count);

         WriteString("\nMem Used: ");
         WriteNumber(bump.offset);

         WriteString("\n--- BumpArena Report ---\n\n");

   }

   return null;
} @jc.nodocs

PagingArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   Page :: struct {
      base:   *u8;
      next:   *Page;
      offset: int;
   }

   PageArenaData :: struct {
      pages:   *Page;
      current: *Page;
   }

   PageNOffset :: size_of(Page);
   Page0Offset :: size_of(PageArenaData) + PageNOffset;

   NewPage :: inline (previous_base: *void) -> *Page {
      mem := MemAcquirePage(.CommitImmediately);

      p := mem.(*Page);
      p.base   = mem;
      p.next   = null;
      p.offset = PageNOffset;

      return p;
   }

   data := arena_data.(*PageArenaData);
   if event == {
      case .Setup;
         // This is a special page that stores both the PageArenaData and Page 0
         mem    := MemAcquirePage(.CommitImmediately);
         header := mem.(*PageArenaData);

         page := (mem + size_of(PageArenaData)).(*Page);
         page.base   = mem;
         page.next   = null;
         page.offset = Page0Offset;

         header.pages   = page;
         header.current = page;

         return header;

      case .Teardown;
         page := data.pages;
         while page != null {
            next := page.next;
            MemReleasePage(page.base, PageSize);
            page = next;
         }

      case .Acquire;
         Assert(data.pages   != null, "jc: PageArena, no pages exist");
         Assert(data.current != null, "jc: PageArena, no current page");

         cur := data.current;
         end := cur.base + cur.offset;
         if new_size == 0
          { return end; }

         ptr  := MemAlignForward(end, align);
         size := new_size + (ptr - end);

         // create or reuse page
         if cur.offset + size > PageSize {
            page: *Page;
            if cur.next != null {
               page = cur.next;
            }
            else {
               cur.next = NewPage(cur.base);
               page = cur.next;
            }


            ptr  = page.base + page.offset;
            size = new_size;
            cur  = page;

            data.current = cur;
         }

         cur.offset += size;
         return ptr;

      case .Reset;
         page := data.pages;
         page.offset = Page0Offset;

         page = page.next;
         while page != null {
            page.offset = PageNOffset;
            page = page.next;
         }

         data.current = data.pages;

      case .Save;
         return data.current.(*void);

      case .Restore;
         if old_ptr == null
          { Panic("jc: PagingArena, restore point was invalid"); }

         // @note(judah): when restoring, we need to reset every page
         // from the savepoint forward
         saved_page := old_ptr.(*Page);
         while saved_page != null {
            if saved_page == data.pages {
               saved_page.offset = Page0Offset;
            }
            else {
               saved_page.offset = PageNOffset;
            }

            saved_page = saved_page.next;
         }

      case .Report;
         pages     := 0;
         allocated := 0;
         in_use    := 0;

         page := data.pages;
         while page != null {
            pages     += 1;
            allocated += PageSize;
            in_use    += page.offset;

            page = page.next;
         }

         WriteString("\n--- PagingArena Report ---");
         WriteString("\nPages: ");
         WriteNumber(pages);

         WriteString("\nMem Used: ");
         WriteNumber(in_use);

         WriteString("\nMem Allocated: ");
         WriteNumber(allocated);
         WriteString(" (PageSize * N Pages)");


         WriteString("\n--- PagingArena Report ---\n\n");
   }

   return null;
} @jc.nodocs

// @note(judah): TempArenaProc just wraps BumpArenaProc, but provides its own backing memory.
// This is needed so we can use TempArenaProc without forcing the user to set everything up.
TempArenaProc :: (
   event: ArenaEvent,
   arena_data: *void,

   new_size: int,   old_size: int, align: int,
   new_ptr:  *void, old_ptr: *void,

   caller: Source_Code_Location
) -> *void {
   // @temp(judah)
   libc :: #library,system "libc";
   malloc :: (size: int) -> *void #c_call #foreign libc;
   free   :: (ptr: *void) #c_call #foreign libc;

   TempMemorySize :: (32 * Kilobyte) + size_of(BumpData);

   if event == {
      // @note(judah): allows the temp arena to initialize itself without relying on another arena to provide its memory.
      case .Setup;
         mcount := AlignForward(TempMemorySize).(int);
         memory := malloc(mcount).(*u8);
         MemZero(memory, mcount);

         bump := memory.(*BumpData);
         bump.memory.data  = memory + size_of(BumpData);
         bump.memory.count = mcount - size_of(BumpData);

         return bump;

      case .Teardown;
         free(arena_data);
   }

   return BumpArenaProc(event, arena_data, new_size, old_size, align, new_ptr, old_ptr, caller);
} @jc.nodocs


#scope_file

JaiAllocatorProc :: (event: Allocator_Mode, requested_size: s64, old_size: s64, old_memory: *void, allocator_data: *void) -> *void {
   CallerFromStackTrace :: () -> Source_Code_Location #expand {
      node := context.stack_trace;
      if node.next != null {
         node = node.next;
      }
      return node.info.location;
   }

   arena := allocator_data.(*Arena);
   if event == {
      case .STARTUP;
         TrySetupArena(arena);
         return null;

      case .ALLOCATE;
         loc := CallerFromStackTrace();
         return ArenaPushSize(requested_size, loc = loc,, arena = arena);

      case .RESIZE;
         loc := CallerFromStackTrace();
         new := ArenaPushSize(requested_size, loc = loc,, arena = arena);
         MemCopy(new, old_memory, old_size);
         return new;

      case .FREE;
         // Arenas don't free single pointers so give this a value that's obviously incorrect
         MemOverwrite(old_memory, old_size, 0xAA);
   }

   return null;
} @jc.nodocs