namespace XFESystem.NativeKernel;

// The inherited boot page tables remain in place. The kernel exclusively owns
// two initially-empty PML4 slots for a 1 TiB demand-committed supervisor arena.
// All edits and TLB invalidations run on the single BSP with interrupts disabled.
public static unsafe class VirtualMemory
{
    public const ulong Base = 0x0000200000000000;
    public const ulong Capacity = 1UL << 40;
    const ulong AddressMask = 0x000ffffffffff000;
    const ulong Present = 1, Writable = 2, NoExecute = 1UL << 63;
    const ulong PageMask = 4095;
    const uint RegionCount = 256;
    const ulong RamRegion = 1, DeviceRegion = 2;
    static ulong* _root;
    static ulong* _regions;
    static ulong _cursor;
    public static bool Ready;

    public static bool Initialize()
    {
        Ready = false;
        if (!PhysicalMemory.Ready || (Native.ReadCr4() & (1UL << 12)) != 0) return false;
        uint* cpu = stackalloc uint[4];
        Native.CpuId(0x80000001, 0, cpu);
        if ((cpu[3] & (1u << 20)) == 0) return false;
        Native.WriteMsr(0xc0000080, Native.ReadMsr(0xc0000080) | (1UL << 11)); // EFER.NXE
        _root = (ulong*)PhysicalMemory.ToVirtual(Native.ReadCr3() & AddressMask);
        _regions = (ulong*)Native.Storage(6);
        if (_root == null || _regions == null) return false;
        uint firstSlot = (uint)((Base >> 39) & 511);
        for (uint i = 0; i < Capacity >> 39; i++) if (_root[firstSlot + i] != 0) return false;
        for (uint i = 0; i < RegionCount * 3; i++) _regions[i] = 0;
        _cursor = Base;
        Ready = true;
        return true;
    }

    static bool RoundSize(ulong bytes, out ulong rounded)
    {
        rounded = 0;
        if (bytes == 0 || bytes > Capacity || bytes > ~0UL - PageMask) return false;
        rounded = (bytes + PageMask) & ~PageMask;
        return true;
    }

    public static ulong Reserve(ulong bytes, ulong alignment)
    {
        return ReserveRegion(bytes, alignment, RamRegion);
    }

    static ulong ReserveRegion(ulong bytes, ulong alignment, ulong kind)
    {
        if (!Ready || !RoundSize(bytes, out ulong rounded)) return 0;
        if (alignment == 0) alignment = 4096;
        if (alignment < 4096 || alignment > Capacity || (alignment & (alignment - 1)) != 0) return 0;
        uint slot = RegionCount;
        for (uint i = 0; i < RegionCount; i++) if (_regions[i * 3 + 2] == 0) { slot = i; break; }
        if (slot == RegionCount || _cursor > ~0UL - (alignment - 1)) return 0;
        ulong address = (_cursor + alignment - 1) & ~(alignment - 1);
        if (address < Base || address > Base + Capacity || rounded > Base + Capacity - address) return 0;
        _regions[slot * 3] = address;
        _regions[slot * 3 + 1] = rounded;
        _regions[slot * 3 + 2] = kind;
        _cursor = address + rounded;
        return address;
    }

    static int RegionFor(ulong address, ulong bytes)
    {
        if (!Ready || (address & PageMask) != 0 || !RoundSize(bytes, out ulong rounded) ||
            address < Base || address >= Base + Capacity || rounded > Base + Capacity - address) return -1;
        for (uint i = 0; i < RegionCount; i++)
        {
            ulong regionBase = _regions[i * 3], length = _regions[i * 3 + 1];
            if (_regions[i * 3 + 2] != 0 && address >= regionBase && address - regionBase < length &&
                rounded <= length - (address - regionBase)) return (int)i;
        }
        return -1;
    }

    static ulong* Leaf(ulong address, bool create, bool splitHuge = false)
    {
        ulong* table = _root;
        for (int shift = 39; shift >= 21; shift -= 9)
        {
            uint index = (uint)((address >> shift) & 511);
            ulong entry = table[index];
            if ((entry & Present) == 0)
            {
                if (!create) return null;
                ulong child = PhysicalMemory.Allocate();
                if (child == 0) return null;
                table[index] = child | Present | Writable;
                Native.MemoryBarrier();
                entry = table[index];
            }
            if ((entry & 0x80) != 0)
            {
                if (!splitHuge || (shift != 30 && shift != 21)) return null;
                ulong child = PhysicalMemory.Allocate();
                if (child == 0) return null;
                ulong* newTable = (ulong*)PhysicalMemory.ToVirtual(child);
                ulong oldPageSize = 1UL << shift, nextSize = oldPageSize >> 9;
                ulong oldBase = (entry & AddressMask) & ~(oldPageSize - 1);
                ulong leafFlags = entry & ~AddressMask;
                if (shift == 30) leafFlags |= entry & 4096; // preserve large-page PAT in new 2 MiB leaves
                else
                {
                    leafFlags &= ~0x80UL;
                    if ((entry & 4096) != 0) leafFlags |= 0x80; // large PAT becomes 4 KiB PAT
                }
                for (uint page = 0; page < 512; page++) newTable[page] = (oldBase + page * nextSize) | leafFlags;
                // Intermediate entries must permit the independently controlled
                // leaf permissions. Original effective NX/RO remain in every leaf.
                table[index] = child | Present | Writable | (entry & 4);
                Native.MemoryBarrier(); Native.InvalidatePage(address);
                entry = table[index];
            }
            table = (ulong*)PhysicalMemory.ToVirtual(entry & AddressMask);
            if (table == null) return null;
        }
        return table + ((address >> 12) & 511);
    }

    public static bool Commit(ulong address, ulong bytes)
    {
        int slot = RegionFor(address, bytes);
        if (slot < 0 || _regions[slot * 3 + 2] != RamRegion || !RoundSize(bytes, out ulong rounded)) return false;
        const ulong NewlyCommitted = 0x200; // PTE software bit, private to one IRQ-excluded operation.
        ulong missing = 0;
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong* leaf = Leaf(address + offset, false);
            if (leaf == null || (*leaf & Present) == 0)
                if (++missing > PhysicalMemory.AvailablePages) return false;
        }
        if (missing == 0) return true;
        ulong committed = 0;
        for (; committed < rounded; committed += 4096)
        {
            ulong* leaf = Leaf(address + committed, true);
            if (leaf == null) break;
            if ((*leaf & Present) != 0) continue; // Existing pages keep their content and protection.
            ulong physical = PhysicalMemory.Allocate();
            if (physical == 0) break;
            *leaf = physical | Present | Writable | NoExecute | NewlyCommitted;
            Native.MemoryBarrier();
            Native.InvalidatePage(address + committed);
        }
        bool success = committed == rounded;
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong current = address + offset;
            ulong* leaf = Leaf(current, false);
            if (leaf != null && (*leaf & NewlyCommitted) != 0)
            {
                if (success) *leaf &= ~NewlyCommitted;
                else
                {
                    if (!PhysicalMemory.Release(*leaf & AddressMask)) Kernel.Fail("VM_COMMIT_ROLLBACK");
                    *leaf = 0; Native.MemoryBarrier(); Native.InvalidatePage(current);
                }
            }
            if (!success) Prune(current);
        }
        return success;
    }

    static bool Empty(ulong* table)
    {
        for (uint i = 0; i < 512; i++) if (table[i] != 0) return false;
        return true;
    }

    static void Prune(ulong address)
    {
        byte* scratch = stackalloc byte[48];
        ulong** parents = (ulong**)scratch;
        ulong* physical = (ulong*)(scratch + 24);
        ulong* table = _root;
        uint depth = 0;
        for (int shift = 39; shift >= 21; shift -= 9)
        {
            ulong* parent = table + ((address >> shift) & 511);
            if ((*parent & Present) == 0 || (*parent & 0x80) != 0) break;
            ulong nextPhysical = *parent & AddressMask;
            ulong* next = (ulong*)PhysicalMemory.ToVirtual(nextPhysical);
            if (next == null) return;
            parents[depth] = parent; physical[depth] = nextPhysical;
            depth++;
            table = next;
        }
        while (depth != 0 && Empty(table))
        {
            depth--;
            // Only allocator-owned tables can be reclaimed.
            if (!PhysicalMemory.Release(physical[depth])) return;
            *parents[depth] = 0;
            Native.MemoryBarrier();
            Native.InvalidatePage(address);
            table = depth == 0 ? _root : (ulong*)PhysicalMemory.ToVirtual(physical[depth - 1]);
        }
    }

    public static bool Decommit(ulong address, ulong bytes)
    {
        int slot = RegionFor(address, bytes);
        if (slot < 0 || _regions[slot * 3 + 2] != RamRegion || !RoundSize(bytes, out ulong rounded)) return false;
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong current = address + offset;
            ulong* leaf = Leaf(current, false);
            if (leaf == null || (*leaf & Present) == 0) { Prune(current); continue; }
            ulong physical = *leaf & AddressMask;
            if (!PhysicalMemory.Release(physical)) return false;
            *leaf = 0;
            Native.MemoryBarrier();
            Native.InvalidatePage(current);
            Prune(current);
        }
        return true;
    }

    public static bool Release(ulong address, ulong bytes)
    {
        int slot = RegionFor(address, bytes);
        if (slot < 0 || _regions[slot * 3 + 2] != RamRegion || !RoundSize(bytes, out ulong rounded) || _regions[slot * 3] != address ||
            _regions[slot * 3 + 1] != rounded || !Decommit(address, rounded)) return false;
        _regions[slot * 3 + 2] = 0;
        if (address + rounded == _cursor) _cursor = address;
        return true;
    }

    public static bool Protect(ulong address, ulong bytes, uint flags)
    {
        if (!Ready || (address & PageMask) != 0 || !RoundSize(bytes, out ulong rounded) ||
            flags == 0 || (flags & ~7u) != 0 || (flags & 1) == 0 || (flags & 6) == 6) return false;
        bool kernelImage = false;
        int slot = RegionFor(address, bytes);
        if (slot >= 0 && _regions[slot * 3 + 2] != RamRegion) return false;
        if (slot < 0)
        {
            ulong start = Native.KernelImageStart(), end = Native.KernelImageEnd();
            if (address < start || address >= end || rounded > end - address) return false;
            kernelImage = true;
        }
        // Preflight the whole range before changing any permission.
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong* leaf = Leaf(address + offset, false, kernelImage);
            if (leaf == null || (*leaf & Present) == 0) return false;
        }
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong* leaf = Leaf(address + offset, false);
            *leaf = (*leaf & AddressMask) | Present | ((flags & 2) != 0 ? Writable : 0) | ((flags & 4) == 0 ? NoExecute : 0);
            Native.MemoryBarrier();
            Native.InvalidatePage(address + offset);
        }
        return true;
    }

    public static ulong MapDevice(ulong physical, ulong bytes)
    {
        return MapDeviceRange(physical, bytes, true);
    }

    public static ulong MapDeviceReadOnly(ulong physical, ulong bytes)
    {
        return MapDeviceRange(physical, bytes, false);
    }

    // MMIO mappings own page-table pages only. Their physical leaves belong to
    // the device, never to PhysicalMemory. No WC claim is made: PCD/PWT select
    // the existing UC PAT entry; a future WC path must configure PAT explicitly.
    static ulong MapDeviceRange(ulong physical, ulong bytes, bool writable)
    {
        if ((physical & PageMask) != 0 || !RoundSize(bytes, out ulong rounded) ||
            physical > AddressMask || rounded > (AddressMask | PageMask) - physical + 1) return 0;
        uint* cpu = stackalloc uint[4];
        Native.CpuId(1, 0, cpu);
        if ((cpu[3] & (1u << 16)) == 0 || ((Native.ReadMsr(0x277) >> 24) & 255) != 0) return 0;
        Native.CpuId(0x80000000, 0, cpu);
        if (cpu[0] < 0x80000008) return 0;
        Native.CpuId(0x80000008, 0, cpu);
        uint physicalBits = cpu[0] & 255;
        if (physicalBits < 32 || physicalBits > 52) return 0;
        ulong physicalLimit = 1UL << (int)physicalBits;
        if (physical >= physicalLimit || rounded > physicalLimit - physical) return 0;
        ulong address = ReserveRegion(rounded, 4096, DeviceRegion);
        if (address == 0) return 0;
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong* leaf = Leaf(address + offset, true);
            if (leaf == null)
            {
                // Also prune the failed address: Leaf may have allocated an
                // intermediate table before the allocator was exhausted.
                if (!UnmapDevice(address, rounded)) Kernel.Fail("VM_DEVICE_ROLLBACK");
                return 0;
            }
            *leaf = (physical + offset) | Present | (writable ? Writable : 0) | NoExecute | 0x18;
            Native.MemoryBarrier(); Native.InvalidatePage(address + offset);
        }
        return address;
    }

    public static bool UnmapDevice(ulong address, ulong bytes)
    {
        int slot = RegionFor(address, bytes);
        if (slot < 0 || _regions[slot * 3 + 2] != DeviceRegion || !RoundSize(bytes, out ulong rounded) ||
            _regions[slot * 3] != address || _regions[slot * 3 + 1] != rounded) return false;
        for (ulong offset = 0; offset < rounded; offset += 4096)
        {
            ulong current = address + offset;
            ulong* leaf = Leaf(current, false);
            if (leaf != null && (*leaf & Present) != 0)
            {
                *leaf = 0;
                Native.MemoryBarrier(); Native.InvalidatePage(current);
            }
            Prune(current);
        }
        _regions[slot * 3 + 2] = 0;
        if (address + rounded == _cursor) _cursor = address;
        return true;
    }

    public static bool SelfTest()
    {
        ulong free = PhysicalMemory.AvailablePages;
        ulong region = Reserve(5 * 4096, 4096);
        if (region == 0 || PhysicalMemory.AvailablePages != free || Leaf(region, false) != null) return false;
        ulong usable = region + 4096;
        if (!Commit(usable, 4096)) return false;
        *(ulong*)usable = 0x1122334455667788;
        if (!Commit(usable, 3 * 4096) || *(ulong*)usable != 0x1122334455667788) return false;
        *(ulong*)usable = 0;
        ulong* memory = (ulong*)usable;
        if (memory[0] != 0 || memory[1535] != 0) return false;
        memory[0] = 0x535953544d564d31;
        memory[1535] = 0x5846455041474532;
        if (memory[0] != 0x535953544d564d31 || memory[1535] != 0x5846455041474532) return false;
        ulong* before = Leaf(region, false), after = Leaf(region + 4 * 4096, false);
        if ((before != null && (*before & Present) != 0) || (after != null && (*after & Present) != 0)) return false;
        if (!Cpu.ProtectionProbe(region, 0) || !Cpu.ProtectionProbe(region + 4 * 4096, 0)) return false;
        if (!Protect(usable, 3 * 4096, 1) || (*Leaf(usable, false) & (Writable | NoExecute)) != NoExecute) return false;
        if (!Cpu.ProtectionProbe(usable, 1)) return false;
        if (!Protect(usable, 3 * 4096, 5) || (*Leaf(usable, false) & (Writable | NoExecute)) != 0) return false;
        if (Protect(usable, 3 * 4096, 7) || !Protect(usable, 3 * 4096, 3)) return false;
        *(byte*)usable = 0xc3; // RET, but the page remains NX and must fault on fetch.
        if (!Cpu.ProtectionProbe(usable, 2)) return false;
        if (!Decommit(usable, 3 * 4096) || !Release(region, 5 * 4096)) return false;
        if (PhysicalMemory.AvailablePages != free) return false;
        // Inspect PTEs without dereferencing the device aperture. This does not
        // send MMIO to a real device or count as a GPU execution self-test.
        ulong device = MapDeviceReadOnly(0xfec00000, 8192);
        if (device == 0 || (*Leaf(device, false) & (AddressMask | Writable | NoExecute | 0x18)) !=
            (0xfec00000UL | NoExecute | 0x18)) return false;
        if (Commit(device, 4096) || Decommit(device, 4096) || Release(device, 8192) ||
            Protect(device, 4096, 3) || UnmapDevice(device, 4096) || UnmapDevice(device + 4096, 4096)) return false;
        if (!UnmapDevice(device, 8192) || UnmapDevice(device, 8192) || PhysicalMemory.AvailablePages != free) return false;
        device = MapDevice(0xfec00000, 4096);
        if (device == 0 || (*Leaf(device, false) & Writable) == 0 || !UnmapDevice(device, 4096)) return false;
        return PhysicalMemory.AvailablePages == free;
    }
}