using Mono.Cecil;
using Mono.Cecil.Cil;

namespace XfeBootPlatformPatch;

/// <summary>Exact IL-only latency policy; no queue, IRQ, wait, GC or mode-setting algorithms are replaced.</summary>
public static class LatencyPolicyPatch
{
    public const string DriversName = "Cosmos.Kernel.Drivers";
    public const string CoreLibName = "System.Private.CoreLib";
    public const string DriversOriginalSha256 = "c65e80e18b9d63a35c2f77384d104624ed9aa69a26db024514847e2997ed32bc";
    public const string OfficialPlugsSha256 = "0118211a282ab2255ab46526cef0d3d0e78311db5804118e39ab42c9d203a2ba";
    public const string OfficialCoreLibSha256 = "64db826d55e769d4d3d639425e0e9f02b847e5ed1be13e4ab0619b2a7eefbec5";
    public const string OriginalSchedulerTimer = "362c3257ae8fcf4dbcaaefb06abdae2989fb079f0e33a9907e0d925045d71c17";
    public const string PatchedSchedulerTimer = "2f751a9978bdc5e83756ee02eaf113d8e8e1d862534e1c06719a0207fc4aa1bd";
    public const string OriginalHotPlug = "dcff5eeeeb7afb3856fcf7773049db5da29dcf27972f9e3af3101e6de0a75594";
    public const string PatchedHotPlug = "b3b037c814637cfddd7c36b2e2ea951d2b49b05f73698c8bf785d382080595bb";
    public const uint SchedulerTimerIntervalMs = 1;
    public const uint XhciPolledDrainIntervalMs = 1;
    public const int SuppressedMonitorLogCalls = 9;
    public sealed record MonitorContract(string Name, string Original, string Patched, int[] LogInstructions);
    public static readonly MonitorContract[] Monitors =
    [
        new("LowLevelMonitor_Acquire", "36d5d38821e2684f4b8036649ac83598b8ce8c4d632504240bb896e874538fb6", "6c80108dd9bddbbb17ba195889745eafe5bac395d0c459b3f165629312b98b64", [..Enumerable.Range(0,5), ..Enumerable.Range(11,5)]),
        new("LowLevelMonitor_Release", "bb8534bad23caf86559ae0db0b5079b7e51661242f4b0b8ea36cd847136db9f7", "4700737d3f2a0aff39342b5484b420d6021652225e9ab764b6cb3a6d23cbc218", [..Enumerable.Range(0,5)]),
        new("LowLevelMonitor_Wait", "ddf1d761ce8d47159e518768c3bf6a786390d1ded79a3aa429c11843bb8e668d", "4acf7905e491ecc8c858463e4a49b739fa413b23855cced3a9ab18aaaf7d731b", [..Enumerable.Range(0,5), ..Enumerable.Range(11,5)]),
        new("LowLevelMonitor_TimedWait", "84ba9b1954ef67fa6ad1aa30e72b2edfa105cb7120569b176bfbfbd2d0845a8c", "6e56962fd277334f026c5d566b385fd0fe734c401c6c41a0fac4ad4283d310ae", [..Enumerable.Range(0,12)]),
        new("LowLevelMonitor_Signal_Release", "94811f783f1b8326228711eee7830e1dcf24714e1cf5fe0937fb662f3be07b0f", "4964a8d107dd450506690b195012a06eb51f760549376b29040e3328801c02e5", [..Enumerable.Range(0,5)])
    ];

    public static MethodDefinition SchedulerTimer(AssemblyDefinition hal) => hal.MainModule.Types.Single(type =>
        type.FullName == "Cosmos.Kernel.HAL.X64.X64PlatformInitializer").Methods.Single(method => method.Name == "StartSchedulerTimer");
    public static MethodDefinition HotPlug(AssemblyDefinition drivers) => drivers.MainModule.Types.Single(type =>
        type.FullName == "Cosmos.Kernel.Drivers.Pci.Bus.Xhci.XhciState").Methods.Single(method => method.Name == "HotPlugMain");
    public static MethodDefinition Monitor(AssemblyDefinition corelib, string name) => corelib.MainModule.Types.Single(type => type.FullName == "Interop")
        .NestedTypes.Single(type => type.FullName == "Interop/Sys").Methods.Single(method => method.Name == name);

    public static void ApplySchedulerTimer(AssemblyDefinition hal)
    {
        var method = SchedulerTimer(hal);
        Require(method, OriginalSchedulerTimer, "original x64 scheduler timer setup");
        var call = method.Body.Instructions.Single(instruction => instruction.Operand is MethodReference target &&
            target.FullName == "System.Void Cosmos.Kernel.Core.X64.Cpu.LocalApic::StartPeriodicTimer(System.UInt32)");
        if (call.Previous.OpCode != OpCodes.Ldarg_1) throw new InvalidOperationException("The LAPIC period is no longer the expected argument.");
        call.Previous.OpCode = OpCodes.Ldc_I4_1; call.Previous.Operand = null;
        Require(method, PatchedSchedulerTimer, "1 ms scheduler timer setup");
    }

    public static void ApplyIo(AssemblyDefinition drivers, AssemblyDefinition corelib)
    {
        VerifyDrivers(drivers);
        var hot = HotPlug(drivers);
        Require(hot, OriginalHotPlug, "original xHCI fallback drain");
        RequireMonitorState(corelib, false);
        var sleep = hot.Body.Instructions.Single(instruction => instruction.Operand is MethodReference target &&
            target.FullName == "System.Void Cosmos.Kernel.HAL.DriverKit.DeviceBinding::Sleep(System.UInt32)");
        if (sleep.Previous.OpCode != OpCodes.Ldc_I4_S || !Equals(sleep.Previous.Operand, (sbyte)20))
            throw new InvalidOperationException("The fallback xHCI period is no longer exactly 20 ms.");
        sleep.Previous.OpCode = OpCodes.Ldc_I4_1; sleep.Previous.Operand = null;
        foreach (var contract in Monitors)
        {
            var method = Monitor(corelib, contract.Name);
            foreach (int index in contract.LogInstructions)
            { method.Body.Instructions[index].OpCode = OpCodes.Nop; method.Body.Instructions[index].Operand = null; }
        }
        VerifyIo(drivers, corelib);
    }

    public static bool OriginalMonitorState(AssemblyDefinition corelib)
    {
        CoreLibIdentity(corelib);
        bool original = Monitors.All(contract => PlatformPatch.Fingerprint(Monitor(corelib, contract.Name)) == contract.Original);
        bool patched = Monitors.All(contract => PlatformPatch.Fingerprint(Monitor(corelib, contract.Name)) == contract.Patched);
        if (!original && !patched) throw new InvalidOperationException("Unknown or partial monitor log IL; refusing to replace waiting semantics.");
        return original;
    }

    public static void VerifyIo(AssemblyDefinition drivers, AssemblyDefinition corelib)
    {
        VerifyDrivers(drivers);
        Require(HotPlug(drivers), PatchedHotPlug, "1 ms no-MSI xHCI drain");
        RequireMonitorState(corelib, true);
    }

    private static void RequireMonitorState(AssemblyDefinition corelib, bool patched)
    {
        CoreLibIdentity(corelib);
        foreach (var contract in Monitors) Require(Monitor(corelib, contract.Name), patched ? contract.Patched : contract.Original, contract.Name);
    }
    private static void CoreLibIdentity(AssemblyDefinition corelib)
    {
        if (corelib.Name.Name != CoreLibName || corelib.Name.Version != new Version(10, 0, 0, 0))
            throw new InvalidOperationException("Only the pinned .NET 10 Cosmos-patched CoreLib is supported.");
    }
    private static void VerifyDrivers(AssemblyDefinition drivers)
    {
        if (drivers.Name.Name != DriversName || drivers.Name.Version != new Version(3, 0, 94, 0))
            throw new InvalidOperationException("Only the pinned Cosmos drivers 3.0.94 assembly is supported.");
    }
    private static void Require(MethodDefinition method, string expected, string description)
    {
        if (!method.HasBody || PlatformPatch.Fingerprint(method) != expected)
            throw new InvalidOperationException("Pinned " + description + " IL differs; no target instruction is guessed.");
    }
}