XFESystem

using System.Runtime.InteropServices;

namespace XFESystem.NativeKernel;

[StructLayout(LayoutKind.Sequential)]
public unsafe struct KernelThread
{
    public ulong State, Id, Reservation, ReservationBytes, StackLow, StackHigh;
    public delegate* unmanaged<void*, uint> Entry;
    public void* Argument;
    public ulong Tls, OwnTls, WakeAt;
    public void* RuntimeThread;
    public InterruptFrame* Interrupt;
    public ulong WaitHandle, ExitResult, InterruptStackTop;
}
[StructLayout(LayoutKind.Sequential)]
public unsafe struct SwitchContext
{
    public ulong Rbx, Rbp, R12, R13, R14, R15, Rsp, Rip, Rflags, FsBase, Rax, Rdx;
    public fixed byte FxState[512];
}
[StructLayout(LayoutKind.Sequential)]
public struct RuntimeLimitedContext { public ulong Ip, Rsp, Rbp, Rax, Rbx, Rdx, R12, R13, R14, R15; }

// Single-BSP scheduler, actual stacks/context switching, no host thread API.
public static unsafe class Threads
{
    const uint Capacity = 32;
    public const ulong InterruptStackBytes = 65536;
    const ulong Runnable = 1, Running = 2, Waiting = 3, Finished = 4;
    static KernelThread* _threads;
    static byte* _contexts;
    static uint _current;
    static ulong _nextId;
    static bool _scheduling;
    public static bool Ready;
    public static ulong Switches, Preemptions;
    public static KernelThread* Current => Ready ? _threads + _current : null;
    static SwitchContext* Context(uint slot) => (SwitchContext*)(_contexts + slot * 1024);

    public static bool Initialize()
    {
        if (Ready || sizeof(KernelThread) != 128 || sizeof(SwitchContext) != 608) return false;
        _threads = (KernelThread*)Native.Storage(7);
        _contexts = (byte*)Native.Storage(8);
        if (_threads == null || _contexts == null) return false;
        for (uint i = 0; i < Capacity; i++) { _threads[i] = new KernelThread(); for (uint j = 0; j < 1024; j++) _contexts[i * 1024 + j] = 0; }
        _current = 0; _nextId = 1;
        _threads[0].State = Running; _threads[0].Id = _nextId++;
        _threads[0].StackHigh = Native.BootStackTop();
        _threads[0].StackLow = _threads[0].StackHigh - 1024 * 1024;
        // The boot stack belongs to Limine. Reserve only a guarded IRQ stack for
        // the BSP thread; new threads place their user and IRQ stacks in one region.
        ulong interruptRegion = VirtualMemory.Reserve(InterruptStackBytes + 8192, 4096);
        if (interruptRegion == 0) return false;
        if (!VirtualMemory.Commit(interruptRegion + 4096, InterruptStackBytes))
        {
            VirtualMemory.Release(interruptRegion, InterruptStackBytes + 8192);
            return false;
        }
        _threads[0].Reservation = interruptRegion;
        _threads[0].ReservationBytes = InterruptStackBytes + 8192;
        _threads[0].InterruptStackTop = interruptRegion + 4096 + InterruptStackBytes;
        if (!Cpu.SetInterruptStack(_threads[0].InterruptStackTop)) return false;
        Ready = true;
        return true;
    }

    public static ulong Create(delegate* unmanaged<void*, uint> entry, void* argument, ulong stackBytes, ulong tls, bool ownTls)
    {
        ulong flags = Native.IrqSave();
        ulong result = CreateLocked(entry, argument, stackBytes, tls, ownTls);
        Native.IrqRestore(flags);
        return result;
    }
    static ulong CreateLocked(delegate* unmanaged<void*, uint> entry, void* argument, ulong stackBytes, ulong tls, bool ownTls)
    {
        if (!Ready || entry == null || stackBytes < 16384 || stackBytes > 32 * 1024 * 1024) return 0;
        Reap();
        uint slot = Capacity;
        for (uint i = 1; i < Capacity; i++) if (_threads[i].State == 0) { slot = i; break; }
        if (slot == Capacity) return 0;
        stackBytes = (stackBytes + 4095) & ~4095UL;
        ulong regionBytes = stackBytes + InterruptStackBytes + 3 * 4096;
        ulong reservation = VirtualMemory.Reserve(regionBytes, 4096);
        ulong interruptLow = reservation + stackBytes + 8192;
        if (reservation == 0 || !VirtualMemory.Commit(reservation + 4096, stackBytes) ||
            !VirtualMemory.Commit(interruptLow, InterruptStackBytes))
        {
            if (reservation != 0) VirtualMemory.Release(reservation, regionBytes);
            return 0;
        }
        SwitchContext* context = Context(slot);
        for (uint i = 0; i < 1024; i++) _contexts[slot * 1024 + i] = 0;
        _threads[slot] = new KernelThread { State = Runnable, Id = _nextId++, Reservation = reservation,
            ReservationBytes = regionBytes, StackLow = reservation + 4096, StackHigh = reservation + 4096 + stackBytes,
            Entry = entry, Argument = argument, Tls = tls, OwnTls = ownTls ? 1UL : 0UL,
            InterruptStackTop = interruptLow + InterruptStackBytes };
        context->Rbx = slot;
        context->Rsp = _threads[slot].StackHigh - 8;
        context->Rip = Native.ThreadTrampolineAddress();
        context->Rflags = 0x202;
        context->FsBase = tls;
        Native.ContextInitialize(context);
        return _threads[slot].Id;
    }

    static void Reap()
    {
        for (uint i = 1; i < Capacity; i++)
        {
            if (i == _current || _threads[i].State != Finished) continue;
            if (!VirtualMemory.Release(_threads[i].Reservation, _threads[i].ReservationBytes)) Kernel.Fail("THREAD_STACK_RELEASE");
            if (_threads[i].OwnTls != 0 && Native.ThreadTlsRelease(_threads[i].Tls) != 1) Kernel.Fail("THREAD_TLS_RELEASE");
            _threads[i].State = 0;
        }
    }
    static void WakeExpired()
    {
        ulong now = MonotonicClock.Nanoseconds();
        for (uint i = 0; i < Capacity; i++)
            if (_threads[i].State == Waiting && _threads[i].WakeAt != ~0UL && now >= _threads[i].WakeAt)
                _threads[i].State = Runnable;
    }
    static uint Choose()
    {
        for (uint step = 1; step <= Capacity; step++)
        {
            uint slot = (_current + step) % Capacity;
            if (_threads[slot].State == Runnable) return slot;
        }
        return Capacity;
    }
    static void Schedule()
    {
        if (!Ready || _scheduling) return;
        _scheduling = true;
        Reap(); WakeExpired();
        uint old = _current;
        if (_threads[old].State == Running) _threads[old].State = Runnable;
        uint next = Choose();
        while (next == Capacity)
        {
            Native.Idle(); WakeExpired(); next = Choose();
        }
        if (next == old) { _threads[old].State = Running; _scheduling = false; return; }
        _threads[next].State = Running;
        _current = next; Switches++;
        if (!Cpu.SetInterruptStack(_threads[next].InterruptStackTop)) Kernel.Fail("THREAD_INTERRUPT_STACK");
        Native.ContextSwitch(Context(old), Context(next));
        _scheduling = false;
        Reap();
    }
    public static void Yield()
    {
        ulong flags = Native.IrqSave(); Schedule(); Native.IrqRestore(flags);
    }
    public static void Sleep(uint milliseconds)
    {
        if (!Ready) return;
        if (milliseconds == 0) { Yield(); return; }
        ulong flags = Native.IrqSave();
        Current->WakeAt = MonotonicClock.Nanoseconds() + (ulong)milliseconds * 1000000;
        Current->WaitHandle = 0; Current->State = Waiting;
        Schedule(); Native.IrqRestore(flags);
    }
    public static void ParkEvent(ulong handle, ulong deadline)
    {
        if (!Ready) return;
        Current->WaitHandle = handle; Current->WakeAt = deadline; Current->State = Waiting;
        Schedule();
        Current->WaitHandle = 0;
    }
    public static void WakeEvent(ulong handle)
    {
        if (!Ready) return;
        for (uint i = 0; i < Capacity; i++)
            if (_threads[i].State == Waiting && (_threads[i].WaitHandle == handle || _threads[i].WaitHandle == ~0UL))
                _threads[i].State = Runnable;
    }
    public static void Timer(InterruptFrame* frame)
    {
        LocalApicTimer.Tick();
        WakeExpired();
        if (!Ready || _scheduling) return;
        uint slot = _current;
        ulong frameAddress = (ulong)frame, top = _threads[slot].InterruptStackTop;
        if (frame == null || frameAddress < top - InterruptStackBytes || frameAddress > top - (ulong)sizeof(InterruptFrame))
            Kernel.Fail("IRQ_STACK_BOUNDS");
        _threads[slot].Interrupt = frame;
        Preemptions++;
        Schedule();
        _threads[slot].Interrupt = null;
    }
    [UnmanagedCallersOnly(EntryPoint = "xfe_kernel_thread_entry")]
    public static void Run(ulong slot)
    {
        if (!Ready || slot >= Capacity || slot != _current) Kernel.Fail("THREAD_ENTRY");
        _scheduling = false;
        Native.IrqRestore(0x202);
        uint result = _threads[slot].Entry(_threads[slot].Argument);
        ulong flags = Native.IrqSave();
        if (Native.ThreadRuntimeExit(_threads[slot].RuntimeThread) != 1) Kernel.Fail("RUNTIME_THREAD_EXIT");
        _threads[slot].ExitResult = result;
        _threads[slot].State = Finished;
        Schedule();
        Native.IrqRestore(flags);
        Kernel.Fail("DEAD_THREAD_RESUMED");
    }
    public static bool StackBounds(void** low, void** high)
    {
        if (!Ready || low == null || high == null) return false;
        *low = (void*)Current->StackLow; *high = (void*)Current->StackHigh;
        return Current->StackLow < Current->StackHigh;
    }
    public static void Attach(void* runtimeThread)
    {
        ulong flags = Native.IrqSave();
        if (!Ready || runtimeThread == null || (Current->RuntimeThread != null && Current->RuntimeThread != runtimeThread)) Kernel.Fail("RUNTIME_ATTACH");
        Current->RuntimeThread = runtimeThread;
        Native.IrqRestore(flags);
    }
    public static int RequestHijack(void* runtimeThread)
    {
        if (!Ready || runtimeThread == null) return -1;
        ulong deadline = MonotonicClock.Nanoseconds() + 5000000000;
        do
        {
            ulong flags = Native.IrqSave();
            int result = -1;
            for (uint slot = 0; slot < Capacity; slot++)
            {
                if (_threads[slot].State == 0 || _threads[slot].State == Finished || _threads[slot].RuntimeThread != runtimeThread) continue;
                RuntimeLimitedContext context;
                InterruptFrame* frame = _threads[slot].Interrupt;
                if (frame != null)
                    context = new RuntimeLimitedContext { Ip = frame->Rip, Rsp = frame->Rsp, Rbp = frame->Rbp,
                        Rax = frame->Rax, Rbx = frame->Rbx, Rdx = frame->Rdx, R12 = frame->R12, R13 = frame->R13, R14 = frame->R14, R15 = frame->R15 };
                else
                {
                    SwitchContext* saved = Context(slot);
                    context = new RuntimeLimitedContext { Ip = saved->Rip, Rsp = saved->Rsp, Rbp = saved->Rbp,
                        Rax = saved->Rax, Rbx = saved->Rbx, Rdx = saved->Rdx, R12 = saved->R12, R13 = saved->R13, R14 = saved->R14, R15 = saved->R15 };
                }
                result = Native.TryHijack(runtimeThread, &context);
                break;
            }
            Native.IrqRestore(flags);
            if (result == 1) return 1;
            if (result < 0) return -1;
            Yield(); // inconvenient context: let the real thread run to a GC poll, then retry
        } while (MonotonicClock.Nanoseconds() < deadline);
        return 0;
    }
}