const std = @import("std");
const builtin = @import("builtin");
const config = @import("pardes_config");

const Allocator = std.mem.Allocator;
const debug_enabled = builtin.mode == .Debug;
/// Three tiers, because the address space differs by four orders of magnitude.
/// `reduced_target` is the browser: a wasm linear memory it grows on demand, so
/// the static reservations are megabytes rather than tens.
///
/// `p4` is ESP32-P4 firmware, and its tier is deliberately ALL FALLBACK. Every
/// capacity here is a `StackFallbackAllocator`'s buffer, which is a static and
/// therefore lands in `.bss` — and on the P4 `.bss`, `.data` and the stack all
/// share ONE 240 KiB chunk of L2MEM at 0x4FF03000, while the heap the fallback
/// allocator hands out is the separate 384 KiB chunk at 0x4FF40000 - the 128 KiB
/// above that measured as L2 cache rather than memory. A
/// megabyte-shaped reservation here would not fit, and every byte that did fit
/// would be taken from the stack's neighbourhood to duplicate memory the heap
/// already has. So the buffers exist only because the type requires one: 4 KiB
/// absorbs the small churn, and everything else spills to the real heap on the
/// first allocation.
const p4 = config.platform == .p4;
const reduced_target = config.platform == .web or builtin.os.tag == .freestanding;
const KiB = 1024;
const MiB = 1024 * KiB;

const capacities = struct {
    const pardes = if (p4) 4 * KiB else if (reduced_target) 8 * MiB else 32 * MiB;
    const frame = if (p4) 4 * KiB else if (reduced_target) 4 * MiB else 16 * MiB;
    // Zero is legal and always spills, which is exactly what an arena for a
    // compiled-out subsystem should do. `StackFallbackAllocator(0).buffer` is
    // `[0]u8`; `get()` inits the FixedBufferAllocator over an empty slice, so
    // `FixedBufferAllocator.alloc` fails every nonzero request and `alloc`
    // falls through to `self.fallback_allocator.rawAlloc`, while `ownsPtr` over
    // an empty range is false for every pointer so `resize`/`remap`/`free`
    // route to the fallback too. See lib/std/heap.zig, StackFallbackAllocator.
    const tree_sitter = if (p4) 0 else if (reduced_target) 4 * MiB else 16 * MiB;
    const image = if (p4) 0 else if (reduced_target) 64 * KiB else 32 * MiB;
    const pdf = if (p4) 0 else if (reduced_target or !config.mupdf) 64 * KiB else 64 * MiB;
};

pub const Allocators = struct {
    pardes: Allocator,
    frame: Allocator,
    lsp: Allocator,
    tree_sitter: Allocator,
    image: Allocator,
    pdf: Allocator,
};

var pardes_fallback: std.heap.StackFallbackAllocator(capacities.pardes) = undefined;
var frame_fallback: std.heap.StackFallbackAllocator(capacities.frame) = undefined;
var tree_sitter_fallback: std.heap.StackFallbackAllocator(capacities.tree_sitter) = undefined;
var image_fallback: std.heap.StackFallbackAllocator(capacities.image) = undefined;
var pdf_fallback: std.heap.StackFallbackAllocator(capacities.pdf) = undefined;

const Debug = std.heap.DebugAllocator(.{});
var pardes_debug: Debug = .init;
var frame_debug: Debug = .init;
var lsp_debug: Debug = .init;
var tree_sitter_debug: Debug = .init;
var image_debug: Debug = .init;
var pdf_debug: Debug = .init;

/// Returns ordinary allocators backed by stdlib fixed-buffer fallbacks. LSP
/// keeps the caller's allocator because its detached workers are concurrent.
pub fn init(fallback: Allocator) Allocators {
    pardes_fallback.fallback_allocator = fallback;
    pardes_fallback.get_called = if (std.debug.runtime_safety) false else {};
    frame_fallback.fallback_allocator = fallback;
    frame_fallback.get_called = if (std.debug.runtime_safety) false else {};
    tree_sitter_fallback.fallback_allocator = fallback;
    tree_sitter_fallback.get_called = if (std.debug.runtime_safety) false else {};
    image_fallback.fallback_allocator = fallback;
    image_fallback.get_called = if (std.debug.runtime_safety) false else {};
    pdf_fallback.fallback_allocator = fallback;
    pdf_fallback.get_called = if (std.debug.runtime_safety) false else {};

    const raw: Allocators = .{
        .pardes = pardes_fallback.get(),
        .frame = frame_fallback.get(),
        .lsp = fallback,
        .tree_sitter = tree_sitter_fallback.get(),
        .image = image_fallback.get(),
        .pdf = pdf_fallback.get(),
    };
    if (!debug_enabled) return raw;

    pardes_debug = .{ .backing_allocator = raw.pardes };
    frame_debug = .{ .backing_allocator = raw.frame };
    lsp_debug = .{ .backing_allocator = raw.lsp };
    tree_sitter_debug = .{ .backing_allocator = raw.tree_sitter };
    image_debug = .{ .backing_allocator = raw.image };
    pdf_debug = .{ .backing_allocator = raw.pdf };
    return .{
        .pardes = pardes_debug.allocator(),
        .frame = frame_debug.allocator(),
        .lsp = lsp_debug.allocator(),
        .tree_sitter = tree_sitter_debug.allocator(),
        .image = image_debug.allocator(),
        .pdf = pdf_debug.allocator(),
    };
}

pub fn deinit() void {
    if (!debug_enabled) return;
    var leaked = pardes_debug.deinit() == .leak;
    leaked = (frame_debug.deinit() == .leak) or leaked;
    leaked = (lsp_debug.deinit() == .leak) or leaked;
    leaked = (tree_sitter_debug.deinit() == .leak) or leaked;
    leaked = (image_debug.deinit() == .leak) or leaked;
    leaked = (pdf_debug.deinit() == .leak) or leaked;
    if (leaked) @panic("allocator leaks detected");
}

test "fixed allocators are separate, spill, and restart" {
    var allocs = init(std.testing.allocator);
    const core = try allocs.pardes.alloc(u8, 32);
    const frame = try allocs.frame.alloc(u8, 32);
    try std.testing.expect(pardes_fallback.fixed_buffer_allocator.ownsPtr(core.ptr));
    try std.testing.expect(frame_fallback.fixed_buffer_allocator.ownsPtr(frame.ptr));
    try std.testing.expect(core.ptr != frame.ptr);

    const spill = try allocs.pardes.alloc(u8, capacities.pardes + 1);
    try std.testing.expect(!pardes_fallback.fixed_buffer_allocator.ownsPtr(spill.ptr));
    allocs.pardes.free(spill);
    allocs.frame.free(frame);
    allocs.pardes.free(core);
    deinit();

    allocs = init(std.testing.allocator);
    defer deinit();
    const restarted = try allocs.pardes.alloc(u8, 32);
    defer allocs.pardes.free(restarted);
    try std.testing.expect(pardes_fallback.fixed_buffer_allocator.ownsPtr(restarted.ptr));
}