#import "std.sx";

// --- Allocator protocol (inline: ctx + fn-ptrs, no vtable indirection) ---

Allocator :: protocol #inline {
    alloc :: (size: s64) -> *void;
    dealloc :: (ptr: *void);
}

// --- CAllocator: stateless allocator that delegates directly to libc ---
//
// Zero-sized struct. Used as the default `context.allocator` at program
// start (see `__sx_default_context` in the codegen). The thunks never
// dereference `self`, so the protocol value's ctx field is `null`.
//
// Unlike GPA, no `init()` is needed — there's nothing to allocate.

CAllocator :: struct {}

impl Allocator for CAllocator {
    alloc :: (self: *CAllocator, size: s64) -> *void {
        return libc_malloc(size);
    }
    dealloc :: (self: *CAllocator, ptr: *void) {
        libc_free(ptr);
    }
}

// --- GPA: general purpose allocator (malloc/free wrapper) ---
//
// `init` returns the GPA by value. Caller binds it to a local; the
// local IS the allocator state, no heap-side allocation for the
// struct itself. `xx gpa` borrows the local under Option 3, so the
// Allocator protocol value's `ctx` points at the local.
//
// Usage:
//   gpa := GPA.init();                                    // GPA
//   push Context.{ allocator = xx gpa, data = null } { ... }
//   print("alloc count: {}\n", gpa.alloc_count);

GPA :: struct {
    alloc_count: s64;

    init :: () -> GPA {
        GPA.{ alloc_count = 0 };
    }
}

impl Allocator for GPA {
    alloc :: (self: *GPA, size: s64) -> *void {
        self.alloc_count += 1;
        return libc_malloc(size);
    }
    dealloc :: (self: *GPA, ptr: *void) {
        self.alloc_count -= 1;
        libc_free(ptr);
    }
}

// --- Arena: multi-chunk bump allocator ---
//
// `init` returns the Arena by value; the caller's local holds the
// state, no heap-side allocation for the struct itself. The arena's
// chunks ARE heap-allocated through the parent allocator, but those
// are owned by `deinit` (or `reset` for the non-first ones).
//
// Usage:
//   gpa   := GPA.init();
//   arena := Arena.init(xx gpa, 4096);                    // Arena
//   push Context.{ allocator = xx arena, data = null } { ... }
//   arena.reset();   // free all chunks except the first
//   arena.deinit();  // free every chunk

ArenaChunk :: struct {
    next: *ArenaChunk;
    cap: s64;
}

Arena :: struct {
    first: *ArenaChunk;
    end_index: s64;
    parent: Allocator;

    add_chunk :: (a: *Arena, min_size: s64) {
        prev_cap := if a.first != null then a.first.cap else 0;
        needed := min_size + 16 + 16;
        len := (prev_cap + needed) * 3 / 2;
        raw := a.parent.alloc(len);
        chunk : *ArenaChunk = xx raw;
        chunk.next = a.first;
        chunk.cap = len;
        a.first = chunk;
        a.end_index = 0;
    }

    init :: (parent_alloc: Allocator, size: s64) -> Arena {
        self : Arena = .{ first = null, end_index = 0, parent = parent_alloc };
        self.add_chunk(size);
        self;
    }

    reset :: (a: *Arena) {
        if a.first != null {
            it := a.first.next;
            while it != null {
                next := it.next;
                a.parent.dealloc(it);
                it = next;
            }
            a.first.next = null;
        }
        a.end_index = 0;
    }

    deinit :: (a: *Arena) {
        it := a.first;
        while it != null {
            next := it.next;
            a.parent.dealloc(it);
            it = next;
        }
        a.first = null;
        a.end_index = 0;
    }
}

impl Allocator for Arena {
    alloc :: (self: *Arena, size: s64) -> *void {
        aligned := (size + 7) & (0 - 8);
        if self.first != null {
            usable := self.first.cap - 16;
            if self.end_index + aligned <= usable {
                buf : [*]u8 = xx self.first;
                ptr := @buf[16 + self.end_index];
                self.end_index = self.end_index + aligned;
                return ptr;
            }
        }
        self.add_chunk(aligned);
        buf : [*]u8 = xx self.first;
        ptr := @buf[16 + self.end_index];
        self.end_index = self.end_index + aligned;
        ptr;
    }
    dealloc :: (self: *Arena, ptr: *void) {}
}

// --- BufAlloc: bump allocator backed by a user-provided slice ---
//
// Usage:
//   stack_buf : [128]u8 = ---;
//   buf := BufAlloc.init(@stack_buf[0], 128);             // *BufAlloc
//   push Context.{ allocator = xx buf, data = null } { ... }
//   buf.reset();

BufAlloc :: struct {
    buf: [*]u8;
    len: s64;
    pos: s64;

    init :: (buf: [*]u8, len: s64) -> *BufAlloc {
        self_size :: size_of(BufAlloc);
        if len < self_size { return null; }
        b : *BufAlloc = xx buf;
        b.buf = @buf[self_size];
        b.len = len - self_size;
        b.pos = 0;
        b;
    }

    reset :: (b: *BufAlloc) {
        b.pos = 0;
    }
}

impl Allocator for BufAlloc {
    alloc :: (self: *BufAlloc, size: s64) -> *void {
        aligned := (size + 7) & (0 - 8);
        if self.pos + aligned > self.len {
            return null;
        }
        ptr := @self.buf[self.pos];
        self.pos = self.pos + aligned;
        ptr;
    }
    dealloc :: (self: *BufAlloc, ptr: *void) {}
}

// --- TrackingAllocator: wraps any Allocator, counts allocs/deallocs ---
//
// Useful for catching leaks during development. Wraps a parent
// Allocator; every call delegates to the parent while updating
// counters. `report()` prints a summary; `leak_count()` returns
// (alloc_count - dealloc_count).
//
// Manual opt-in pattern (compiler auto-wrap lands in Phase 5):
//
//   tracker := TrackingAllocator.init(context.allocator);   // TrackingAllocator
//   push Context.{ allocator = xx tracker, data = null } {
//       // ... user code allocates via tracker → delegates to the
//       //     original context.allocator (libc-backed by default) ...
//   }
//   tracker.report();
//   if tracker.leak_count() != 0 { return 1; }
//
// Limitations under the current 2-method Allocator protocol:
// dealloc(ptr) provides no size info, so bytes_outstanding /
// peak_bytes cannot be tracked accurately. Only alloc count and
// total bytes allocated are recorded. Phase 4's size-aware
// dealloc(ptr, size, align) unlocks full byte tracking.

TrackingAllocator :: struct {
    parent: Allocator;
    alloc_count: s64;
    dealloc_count: s64;
    total_alloc_bytes: s64;

    init :: (parent_alloc: Allocator) -> TrackingAllocator {
        TrackingAllocator.{
            parent = parent_alloc,
            alloc_count = 0,
            dealloc_count = 0,
            total_alloc_bytes = 0,
        };
    }

    leak_count :: (t: *TrackingAllocator) -> s64 {
        t.alloc_count - t.dealloc_count;
    }

    report :: (t: *TrackingAllocator) {
        print("TrackingAllocator: allocs={} deallocs={} outstanding={} total_alloc_bytes={}\n",
              t.alloc_count, t.dealloc_count, t.leak_count(), t.total_alloc_bytes);
    }
}

impl Allocator for TrackingAllocator {
    alloc :: (self: *TrackingAllocator, size: s64) -> *void {
        ptr := self.parent.alloc(size);
        if ptr != null {
            self.alloc_count += 1;
            self.total_alloc_bytes += size;
        }
        ptr;
    }
    dealloc :: (self: *TrackingAllocator, ptr: *void) {
        self.parent.dealloc(ptr);
        self.dealloc_count += 1;
    }
}