namespace XFESystem.NativeKernel;

// Single-BSP bootstrap allocator. No CLR heap, third-party allocator or bootloader
// reclaimable region is used. Two bitmaps distinguish free and owned pages.
public static unsafe class PhysicalMemory
{
    public const ulong PageSize = 4096;
    public const ulong MaximumPhysical = 8UL * 1024 * 1024 * 1024;
    const ulong MaximumPages = MaximumPhysical / PageSize;
    const ulong BitmapBytes = MaximumPages / 8;
    static byte* _free;
    static byte* _owned;
    static byte* _starts;
    static byte* _ends;
    static ulong _next;
    public static ulong Hhdm, AvailablePages, AllocatedPages, IgnoredUsableBytes, TotalPages;
    public static bool Ready;

    public static bool Initialize(NativeBootInfo* boot)
    {
        Ready = false;
        if (boot == null || boot->MemoryEntries == null || boot->MemoryEntryCount == 0 ||
            boot->MemoryEntryCount > 4096 || boot->Hhdm == 0) return false;
        _free = (byte*)Native.Storage(4);
        _owned = (byte*)Native.Storage(5);
        _starts = (byte*)Native.Storage(12);
        _ends = (byte*)Native.Storage(13);
        if (_free == null || _owned == null || _starts == null || _ends == null) return false;
        for (ulong i = 0; i < BitmapBytes; i++) { _free[i] = 0; _owned[i] = 0; _starts[i] = 0; _ends[i] = 0; }
        Hhdm = boot->Hhdm;
        AvailablePages = AllocatedPages = IgnoredUsableBytes = _next = 0;
        for (ulong i = 0; i < boot->MemoryEntryCount; i++)
        {
            MemoryMapEntry* entry = boot->MemoryEntries[i];
            if (entry == null || entry->Base > ~0UL - entry->Length) return false;
            ulong end = entry->Base + entry->Length;
            // Reject overlapping map entries rather than count or allocate a page twice.
            for (ulong j = 0; j < i; j++)
            {
                MemoryMapEntry* prior = boot->MemoryEntries[j];
                if (prior->Length != 0 && entry->Length != 0 &&
                    entry->Base < prior->Base + prior->Length && prior->Base < end) return false;
            }
            if (entry->Kind != 0 || entry->Length == 0) continue;
            ulong start = entry->Base;
            if (start < 1024 * 1024) start = 1024 * 1024;
            if (end > MaximumPhysical)
            {
                ulong ignoredStart = start > MaximumPhysical ? start : MaximumPhysical;
                if (end > ignoredStart) IgnoredUsableBytes += end - ignoredStart;
                end = MaximumPhysical;
            }
            if (start >= end || start > ~0UL - (PageSize - 1)) continue;
            start = (start + PageSize - 1) & ~(PageSize - 1);
            end &= ~(PageSize - 1);
            for (ulong page = start / PageSize; page < end / PageSize; page++)
            {
                _free[page >> 3] |= (byte)(1 << (int)(page & 7));
                AvailablePages++;
            }
        }
        if (AvailablePages == 0 || Hhdm > ~0UL - MaximumPhysical) return false;
        TotalPages = AvailablePages;
        Ready = true;
        return true;
    }

    public static ulong Allocate(uint count = 1)
    {
        if (!Ready || count == 0 || count > AvailablePages) return 0;
        ulong run = 0, start = 0;
        // Linear contiguous first-fit, wrapping only between complete attempts.
        for (uint attempt = 0; attempt < 2; attempt++)
        {
            ulong begin = attempt == 0 ? _next : 0;
            ulong limit = attempt == 0 ? MaximumPages : _next;
            run = 0;
            for (ulong page = begin; page < limit; page++)
            {
                byte mask = (byte)(1 << (int)(page & 7));
                if ((_free[page >> 3] & mask) == 0) { run = 0; continue; }
                if (run++ == 0) start = page;
                if (run != count) continue;
                for (ulong p = start; p < start + count; p++)
                {
                    byte bit = (byte)(1 << (int)(p & 7));
                    _free[p >> 3] &= (byte)~bit;
                    _owned[p >> 3] |= bit;
                }
                AvailablePages -= count;
                AllocatedPages += count;
                _starts[start >> 3] |= (byte)(1 << (int)(start & 7));
                ulong last = start + count - 1;
                _ends[last >> 3] |= (byte)(1 << (int)(last & 7));
                _next = start + count;
                byte* bytes = (byte*)ToVirtual(start * PageSize);
                for (ulong i = 0; i < count * PageSize; i++) bytes[i] = 0;
                return start * PageSize;
            }
        }
        return 0;
    }

    public static bool Release(ulong physical, uint count = 1)
    {
        if (!Ready || count == 0 || physical == 0 || (physical & (PageSize - 1)) != 0 ||
            physical >= MaximumPhysical || count > (MaximumPhysical - physical) / PageSize) return false;
        ulong start = physical / PageSize;
        ulong last = start + count - 1;
        if ((_starts[start >> 3] & (1 << (int)(start & 7))) == 0 ||
            (_ends[last >> 3] & (1 << (int)(last & 7))) == 0) return false;
        for (ulong p = start; p < start + count; p++)
        {
            byte bit = (byte)(1 << (int)(p & 7));
            if ((_owned[p >> 3] & bit) == 0 || (p != start && (_starts[p >> 3] & bit) != 0) ||
                (p != last && (_ends[p >> 3] & bit) != 0)) return false;
        }
        for (ulong p = start; p < start + count; p++)
        {
            byte bit = (byte)(1 << (int)(p & 7));
            _owned[p >> 3] &= (byte)~bit;
            _free[p >> 3] |= bit;
        }
        _starts[start >> 3] &= (byte)~(1 << (int)(start & 7));
        _ends[last >> 3] &= (byte)~(1 << (int)(last & 7));
        AvailablePages += count;
        AllocatedPages -= count;
        if (start < _next) _next = start;
        return true;
    }

    public static ulong ToVirtual(ulong physical)
    {
        if (!Ready || physical >= MaximumPhysical || Hhdm > ~0UL - physical) return 0;
        return Hhdm + physical;
    }

    public static bool SelfTest()
    {
        ulong available = AvailablePages;
        ulong first = Allocate(), second = Allocate();
        if (first == 0 || second == 0 || first == second || (first & 4095) != 0) return false;
        ulong* value = (ulong*)ToVirtual(first);
        for (uint i = 0; i < PageSize / 8; i++) if (value[i] != 0) return false;
        value[0] = 0x584645535953544d;
        value[511] = 0xcafebabef00d1234;
        if (value[0] != 0x584645535953544d || value[511] != 0xcafebabef00d1234) return false;
        if (Release(first, 2) || !Release(first) || Release(first) || Release(second + 1)) return false;
        ulong recycled = Allocate();
        if (recycled != first || *(ulong*)ToVirtual(recycled) != 0) return false;
        if (!Release(recycled) || !Release(second)) return false;
        ulong contiguous = Allocate(3);
        if (contiguous == 0 || Release(contiguous, 2) || Release(contiguous + 4096, 2) || !Release(contiguous, 3)) return false;
        return AvailablePages == available && AllocatedPages == 0;
    }
}