//! Physical memory manager: a bitmap frame allocator. It hands out and reclaims
//! 4 KiB physical frames — the primitive every later memory feature (page
//! tables, the heap) is built on top of.
//!
//! This is generic kernel code: it works on the neutral `boot_handoff.MemoryRegion`
//! array the loader hands over (see docs/memory-map.md), so it carries no UEFI
//! and nothing architecture-specific beyond the 4 KiB page.

const std = @import("std");
const boot_handoff = @import("boot-handoff");
const abi = @import("abi");

const page_size = abi.page_size;

/// One bit per frame, covering physical RAM from 0 up to the highest usable
/// address: 1 = used/unavailable, 0 = free. The bitmap itself lives in a frame
/// we carve out of usable memory during init.
var bitmap: []u8 = &.{};
var total_frames: usize = 0;
var used_frames: usize = 0;
/// Where the next allocation scan begins, so we don't rescan from frame 0 every
/// time. Pulled back on free() so reclaimed low frames get reused.
var next_hint: usize = 0;

pub const Stats = struct {
    total_frames: usize,
    used_frames: usize,
    free_frames: usize,
};

pub fn stats() Stats {
    return .{
        .total_frames = total_frames,
        .used_frames = used_frames,
        .free_frames = total_frames - used_frames,
    };
}

inline fn bit(frame: usize) u3 {
    return @intCast(frame & 7);
}
inline fn isUsed(frame: usize) bool {
    return (bitmap[frame >> 3] >> bit(frame)) & 1 != 0;
}
inline fn setUsed(frame: usize) void {
    bitmap[frame >> 3] |= @as(u8, 1) << bit(frame);
}
inline fn setFree(frame: usize) void {
    bitmap[frame >> 3] &= ~(@as(u8, 1) << bit(frame));
}

fn regions(map: boot_handoff.MemoryMap) []const boot_handoff.MemoryRegion {
    return @as([*]const boot_handoff.MemoryRegion, @ptrFromInt(boot_handoff.physicalToVirtual(map.regions)))[0..map.len];
}

/// Build the allocator from the loader's memory map. Reaches physical memory
/// (the region array, the bitmap's own storage) through the physmap, which the
/// loader's bootstrap tables already provide — so this works before the kernel
/// installs its own tables. Invariant: the bitmap lands in the first usable
/// region (lowest address), which must sit under the bootstrap physmap's reach
/// (4 GiB); it always does, as both this and the page-table allocator scan from
/// low addresses up.
pub fn init(map: boot_handoff.MemoryMap) void {
    const regs = regions(map);

    // 1. Size the bitmap to cover every frame up to the highest RAM address —
    //    including reserved RAM, so those frames are trackable (e.g. to free the
    //    boot buffers later). Only MMIO (device address space) is excluded.
    //    Everything starts unallocatable; usable regions are freed below.
    var highest: u64 = 0;
    for (regs) |r| {
        if (r.kind == .mmio) continue;
        const end = r.base + r.pages * page_size;
        if (end > highest) highest = end;
    }
    total_frames = @intCast(highest / page_size);
    if (total_frames == 0) @panic("pmm: no usable memory");
    const bitmap_bytes = (total_frames + 7) / 8;
    const bitmap_pages = (bitmap_bytes + page_size - 1) / page_size;

    // 2. Park the bitmap in the first usable region large enough to hold it.
    //    Start at least one page in, so we never place it on frame 0 (which is
    //    kept reserved as the "none" address, and is an awkward pointer besides).
    var storage: ?u64 = null;
    for (regs) |r| {
        if (r.kind != .usable) continue;
        const base = if (r.base == 0) page_size else r.base;
        const skipped = (base - r.base) / page_size;
        if (r.pages - skipped >= bitmap_pages) {
            storage = base;
            break;
        }
    }
    const bitmap_base = storage orelse @panic("pmm: no region large enough for the frame bitmap");
    bitmap = @as([*]u8, @ptrFromInt(boot_handoff.physicalToVirtual(bitmap_base)))[0..bitmap_bytes];

    // 3. Start with everything marked used, then free the usable regions. Doing
    //    it this way means every gap, reserved span and MMIO hole is unallocatable
    //    by default — we only ever hand back memory the firmware called usable.
    @memset(bitmap, 0xff);
    used_frames = total_frames;
    for (regs) |r| {
        if (r.kind != .usable) continue;
        var f: usize = @intCast(r.base / page_size);
        const end = f + @as(usize, @intCast(r.pages));
        while (f < end and f < total_frames) : (f += 1) {
            setFree(f);
            used_frames -= 1;
        }
    }

    // 4. Take back the frames the bitmap occupies, and frame 0 — so a 0 result
    //    stays reserved to mean "no frame".
    reserve(bitmap_base, bitmap_pages);
    reserve(0, 1);
}

/// Mark `count` frames from physical `base` as used, counting only those that
/// were actually free.
fn reserve(base: u64, count: usize) void {
    var f: usize = @intCast(base / page_size);
    const end = f + count;
    while (f < end and f < total_frames) : (f += 1) {
        if (!isUsed(f)) {
            setUsed(f);
            used_frames += 1;
        }
    }
}

/// Allocate one physical frame, or null if none are free. The address is
/// page-aligned; the frame's contents are undefined.
pub fn alloc() ?u64 {
    var scanned: usize = 0;
    var f = next_hint;
    while (scanned < total_frames) : (scanned += 1) {
        if (f >= total_frames) f = 0;
        if (!isUsed(f)) {
            setUsed(f);
            used_frames += 1;
            next_hint = f + 1;
            return @as(u64, f) * page_size;
        }
        f += 1;
    }
    return null; // out of physical memory
}

/// Allocate one free frame whose physical address is below `limit`, or null if
/// none is free down there. The AP trampoline needs this: an x86 STARTUP IPI vectors
/// a waking core to physical `vector << 12`, and `vector` is a byte — so the
/// trampoline must live under 1 MiB. A short linear scan of the low frames; only run
/// a handful of times at boot, so it needn't be fast.
pub fn allocBelow(limit: u64) ?u64 {
    const cap = @min(total_frames, @as(usize, @intCast(limit / page_size)));
    var f: usize = 1; // frame 0 stays reserved as the "none" address
    while (f < cap) : (f += 1) {
        if (!isUsed(f)) {
            setUsed(f);
            used_frames += 1;
            return @as(u64, f) * page_size;
        }
    }
    return null;
}

/// Return a frame obtained from alloc() to the pool. Bogus or double frees are
/// ignored rather than corrupting the count.
pub fn free(address: u64) void {
    const f: usize = @intCast(address / page_size);
    if (f >= total_frames or !isUsed(f)) return;
    setFree(f);
    used_frames -= 1;
    if (f < next_hint) next_hint = f;
}

/// Allocate `count` physically **contiguous** frames whose highest byte is below
/// `max_phys` (pass `~0` for no limit; use a real limit for DMA engines with 32-bit
/// addressing). Returns the physical base, or null if no free run of that size fits.
/// A DMA descriptor ring needs contiguity, a known physical address, and pinning —
/// none of which one-frame `alloc` gives. Linear scan for a run of clear bits: fine
/// for the small rings DMA needs; a buddy allocator is a later optimisation. The
/// frames are freed individually with `free`, so there is no bespoke free path.
pub fn allocContiguous(count: usize, max_phys: u64) ?u64 {
    if (count == 0) return null;
    const limit: usize = @intCast(@min(@as(u64, total_frames), max_phys / page_size));
    var start: usize = 1; // frame 0 stays reserved as the "none" address
    while (start + count <= limit) {
        if (isUsed(start)) {
            start += 1;
            continue;
        }
        var run: usize = 0;
        while (run < count and !isUsed(start + run)) : (run += 1) {}
        if (run == count) {
            for (0..count) |i| setUsed(start + i);
            used_frames += count;
            return @as(u64, start) * page_size;
        }
        start += run + 1; // the frame at start+run is used; skip past it
    }
    return null; // no contiguous run of `count` frames below max_phys
}