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

const board = config.platform == .esp32p4;
const reduced_target = builtin.os.tag == .freestanding;
const KiB = 1024;
const MiB = 1024 * KiB;

pub const limits = struct {
    pub const max_file_bytes = 256 * MiB;
    pub const max_stream_bytes = 4 * MiB;
    // P4: 384 KiB heap; static buffers and stack share a separate 240 KiB region.
    pub const board_heap_bytes = 384 * KiB;
    pub const effect_cap = if (board) 128 else 4096;
    pub const pending_write_cap: usize = if (board) 0 else 4 * MiB;
    pub const wrap_rows = if (board) 128 else 256;
    pub const undo_max = if (board) 16 else 256;
    pub const message_log = if (board) 16 else 128;
    /// ranges one selection holds; helix has no limit, and a `%s` over a
    /// file's words wants hundreds
    pub const selections = if (board) 64 else 1024;
    /// /log's ring, which records whether or not anyone reads it.
    pub const log_bytes: usize = if (board) 4 * KiB else 64 * KiB;
    pub const max_tag_tail: usize = if (board) 512 else 4096;
    pub const host_path_cap: usize = if (board) 0 else 4095;
    /// -Dembed-sources: the editor serves its own sources under /src.
    pub const embedded_sources = config.embed_sources;
    pub const hexdump_row_bytes: u32 = if (board) 8 else 16;

    pub const arena = struct {
        pub const pardes = if (board) 4 * KiB else if (reduced_target) 8 * MiB else 32 * MiB;
        pub const frame = if (board) 4 * KiB else if (reduced_target) 4 * MiB else 16 * MiB;
        pub const tree_sitter = if (board) 0 else if (reduced_target) 4 * MiB else 16 * MiB;
        pub const image = if (board) 0 else if (reduced_target) 64 * KiB else 32 * MiB;
        pub const pdf = if (board) 0 else if (reduced_target or !config.mupdf) 64 * KiB else 64 * MiB;
    };
};

const Allocator = std.mem.Allocator;
const debug_enabled = builtin.mode == .Debug;

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

/// `std.heap.StackFallbackAllocator` with the buffer taken through its
/// lock-free interface. The core's allocator is used from two threads: a 9P
/// connection task answers a request while the editor's own thread is out
/// in a syscall that allocates (`readFileLimit`, `grep`; see `pardes.turn`).
/// The buffer is a separate zeroed global, not a field: an `undefined` one is
/// written into a Debug binary as 0xAA filler, 160 MB of file for the five
/// arenas, where a zeroed one is .bss and costs nothing.
const SharedStackFallback = struct {
    buffer: []u8,
    fallback_allocator: Allocator = undefined,
    fixed: std.heap.FixedBufferAllocator = undefined,

    fn get(self: *@This()) Allocator {
        self.fixed = .init(self.buffer);
        return .{ .ptr = self, .vtable = &.{ .alloc = alloc, .resize = resize, .remap = remap, .free = free } };
    }

    fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
        const self: *@This() = @ptrCast(@alignCast(ctx));
        return self.fixed.threadSafeAllocator().rawAlloc(len, alignment, ra) orelse
            self.fallback_allocator.rawAlloc(len, alignment, ra);
    }

    fn resize(ctx: *anyopaque, buf: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) bool {
        const self: *@This() = @ptrCast(@alignCast(ctx));
        if (self.fixed.ownsPtr(buf.ptr)) return self.fixed.threadSafeAllocator().rawResize(buf, alignment, new_len, ra);
        return self.fallback_allocator.rawResize(buf, alignment, new_len, ra);
    }

    fn remap(ctx: *anyopaque, buf: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) ?[*]u8 {
        const self: *@This() = @ptrCast(@alignCast(ctx));
        if (self.fixed.ownsPtr(buf.ptr)) return self.fixed.threadSafeAllocator().rawRemap(buf, alignment, new_len, ra);
        return self.fallback_allocator.rawRemap(buf, alignment, new_len, ra);
    }

    fn free(ctx: *anyopaque, buf: []u8, alignment: std.mem.Alignment, ra: usize) void {
        const self: *@This() = @ptrCast(@alignCast(ctx));
        if (self.fixed.ownsPtr(buf.ptr)) return self.fixed.threadSafeAllocator().rawFree(buf, alignment, ra);
        return self.fallback_allocator.rawFree(buf, alignment, ra);
    }
};

var pardes_buffer: [limits.arena.pardes]u8 = @splat(0);
var frame_buffer: [limits.arena.frame]u8 = @splat(0);
var tree_sitter_buffer: [limits.arena.tree_sitter]u8 = @splat(0);
var image_buffer: [limits.arena.image]u8 = @splat(0);
var pdf_buffer: [limits.arena.pdf]u8 = @splat(0);
var pardes_fallback: SharedStackFallback = .{ .buffer = &pardes_buffer };
var frame_fallback: SharedStackFallback = .{ .buffer = &frame_buffer };
var tree_sitter_fallback: SharedStackFallback = .{ .buffer = &tree_sitter_buffer };
var image_fallback: SharedStackFallback = .{ .buffer = &image_buffer };
var pdf_fallback: SharedStackFallback = .{ .buffer = &pdf_buffer };

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;

// One session at a time; free its allocations before deinit. Concurrent LSP workers use the caller's allocator.
pub fn init(fallback: Allocator) Allocators {
    // A Debug build puts each allocator under a DebugAllocator of its own,
    // which reports that subsystem's leaks and bad frees (deinit panics on a
    // leak). Beneath that goes the page allocator rather than `fallback`: in
    // a Debug build `fallback` is std.process.Init's gpa, a DebugAllocator
    // already, and a second one checked every block twice.
    const backing = if (debug_enabled) std.heap.page_allocator else fallback;
    pardes_fallback.fallback_allocator = backing;
    frame_fallback.fallback_allocator = backing;
    tree_sitter_fallback.fallback_allocator = backing;
    image_fallback.fallback_allocator = backing;
    pdf_fallback.fallback_allocator = backing;

    const raw: Allocators = .{
        .pardes = pardes_fallback.get(),
        .frame = frame_fallback.get(),
        .lsp = backing,
        .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.ownsPtr(core.ptr));
    try std.testing.expect(frame_fallback.fixed.ownsPtr(frame.ptr));
    try std.testing.expect(core.ptr != frame.ptr);

    const spill = try allocs.pardes.alloc(u8, limits.arena.pardes + 1);
    try std.testing.expect(!pardes_fallback.fixed.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.ownsPtr(restarted.ptr));
}

test "memory limits preserve desktop capacities" {
    if (board or reduced_target) return error.SkipZigTest;
    try std.testing.expectEqual(4096, limits.effect_cap);
    try std.testing.expectEqual(@as(usize, 4 * MiB), limits.pending_write_cap);
    try std.testing.expectEqual(256, limits.wrap_rows);
    try std.testing.expectEqual(256, limits.undo_max);
    try std.testing.expectEqual(128, limits.message_log);
    try std.testing.expectEqual(1024, limits.selections);
    try std.testing.expectEqual(@as(usize, 4096), limits.max_tag_tail);
    try std.testing.expectEqual(@as(usize, 4095), limits.host_path_cap);
    try std.testing.expectEqual(config.embed_sources, limits.embedded_sources);
    try std.testing.expectEqual(@as(u32, 16), limits.hexdump_row_bytes);
    try std.testing.expectEqual(32 * MiB, limits.arena.pardes);
    try std.testing.expectEqual(16 * MiB, limits.arena.frame);
    try std.testing.expectEqual(16 * MiB, limits.arena.tree_sitter);
    try std.testing.expectEqual(32 * MiB, limits.arena.image);
    try std.testing.expectEqual(if (config.mupdf) 64 * MiB else 64 * KiB, limits.arena.pdf);
}