XFESystem

using System.Globalization;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Text;
using Mono.Cecil;
using Mono.Cecil.Cil;

namespace XfeBootPlatformPatch;

/// <summary>The executing machine has one boot processor; application processors are not started.</summary>
public static class PlatformPatch
{
    public const string PackageVersion = "3.0.94";
    public const string CoreName = "Cosmos.Kernel.Core.X64";
    public const string HalName = "Cosmos.Kernel.HAL.X64";
    public const string CoreOriginalSha256 = "95359dc648a01a2e31ae5fbb513efa3fd796acebd55b69c244c9d4ab988578c0";
    public const string HalOriginalSha256 = "6267612213cbf40a3b20bf56a763652f9075363fb2561d71036d4fe23b164efd";
    public const string PhysicalExport = "XfeRuntimePhysicalCpuCount";
    public const string SchedulerExport = "XfeRuntimeSchedulerCpuCount";
    private const string OriginalTimer = "97693c4fa91997732b61aadc543425be326814a3a7fe5ecbd6fe14316dd70871";
    private const string PatchedTimer = "948a18f6c44247740881c219567b4b1830bfc5dbeb8b133350412bbd70335500";
    private const string OriginalCount = "76314847f293220be85f14b146ed36f9447ebf583898dff04fb0bc21b91600dc";
    private const string PatchedCount = "b9bd48603446583f1db361633209435df227265e38f3883d1a04f2c73bfe01ba";
    private const string OriginalDispatch = "8c41e004b866acbece6ab3a79f1fe856562b7900732a2ba254591233c90f08e5";
    private const string PatchedDispatch = "5e50d316eafcc7d1a4f11aee6dbb0c560a230d25cf1bbbf1decdd517cc844d37";
    private const string PhysicalMethod = "XfePhysicalCpuCountV1";
    private const string SchedulerMethod = "XfeSchedulerCpuCountV1";

    public static void Apply(AssemblyDefinition core, AssemblyDefinition hal)
    {
        Versions(core, hal);
        var timer = Timer(core);
        var count = Count(hal);
        var dispatch = Dispatch(core);
        if (Fingerprint(timer) != OriginalTimer || Fingerprint(count) != OriginalCount || Fingerprint(dispatch) != OriginalDispatch)
            throw new InvalidOperationException("Pinned platform IL changed; refusing to guess a patch location.");
        if (Fingerprint(LatencyPolicyPatch.SchedulerTimer(hal)) != LatencyPolicyPatch.OriginalSchedulerTimer)
            throw new InvalidOperationException("Pinned LAPIC scheduler setup changed.");
        if (hal.MainModule.Types.SelectMany(type => type.Methods).Any(method =>
                method.Name is PhysicalMethod or SchedulerMethod || method.CustomAttributes.Any(attribute =>
                    attribute.AttributeType.FullName == typeof(UnmanagedCallersOnlyAttribute).FullName &&
                    attribute.Fields.Any(field => field.Name == "EntryPoint" && field.Argument.Value is string name &&
                        name is PhysicalExport or SchedulerExport))))
            throw new InvalidOperationException("A platform diagnostic export name is already used.");

        // Retain the exact ACPI getter reference for a read-only physical-topology diagnostic.
        var physical = (MethodReference)count.Body.Instructions[0].Operand;
        count.Body.Instructions.Clear();
        var countIl = count.Body.GetILProcessor();
        countIl.Append(countIl.Create(OpCodes.Ldc_I4_1));
        countIl.Append(countIl.Create(OpCodes.Ret));
        AddExport(hal, PhysicalMethod, PhysicalExport, physical);
        AddExport(hal, SchedulerMethod, SchedulerExport, null);

        var cpu = timer.Body.Instructions.Single(instruction => instruction.Offset == 0x0c);
        cpu.OpCode = OpCodes.Ldc_I4_0;
        cpu.Operand = null;
        var schedulerArguments = timer.Body.Instructions.Single(instruction => instruction.Offset == 0x105);
        var validStack = timer.Body.Instructions.Single(instruction => instruction.Offset == 0x35);
        validStack.Operand = schedulerArguments;
        // Remove the complete normal first-five/every-100-ticks output block, including
        // params object boxing, hexadecimal RSP, and newline. Error diagnostics stay intact.
        var removed = timer.Body.Instructions.Where(instruction => instruction.Offset >= 0x95 && instruction.Offset < 0x105).ToArray();
        foreach (var instruction in removed) timer.Body.Instructions.Remove(instruction);
        // Device-IRQ exit uses the same logical index as the timer. A hardware
        // APIC ID is still used by unchanged routing code, never as this array index.
        var deviceCpu = dispatch.Body.Instructions.Single(instruction => instruction.Offset == 0x73);
        deviceCpu.OpCode = OpCodes.Ldc_I4_0;
        deviceCpu.Operand = null;
        LatencyPolicyPatch.ApplySchedulerTimer(hal);
        Verify(core, hal);
    }

    public static void Verify(AssemblyDefinition core, AssemblyDefinition hal)
    {
        Versions(core, hal);
        var timer = Timer(core);
        var count = Count(hal);
        var dispatch = Dispatch(core);
        if (Fingerprint(timer) != PatchedTimer || Fingerprint(count) != PatchedCount || Fingerprint(dispatch) != PatchedDispatch)
            throw new InvalidOperationException("The installed platform patch differs from its fixed IL contract.");
        if (Fingerprint(LatencyPolicyPatch.SchedulerTimer(hal)) != LatencyPolicyPatch.PatchedSchedulerTimer)
            throw new InvalidOperationException("The installed LAPIC scheduler is not the exact 1 ms contract.");
        var instructions = timer.Body.Instructions;
        if (instructions.Count(instruction => instruction.Operand is MethodReference method && method.Name == "OnTimerInterrupt" &&
                method.DeclaringType.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerManager") != 1 ||
            instructions.Any(instruction => instruction.Operand is string text && text == "[LAPIC] Timer tick ") ||
            instructions.Any(instruction => instruction.Operand is MethodReference method && method.Name is "GetId" or "SendEOI") ||
            instructions.Count(instruction => instruction.Operand is string text && text == "[LAPIC] ERROR: Invalid RSP at tick ") != 1)
            throw new InvalidOperationException("The timer's tick, error, logical-CPU or EOI contract changed.");
        if (dispatch.Body.Instructions.Count(instruction => instruction.Operand is MethodReference method && method.Name == "ReschedulePendingFromIrq" &&
                method.DeclaringType.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerManager") != 1 ||
            dispatch.Body.Instructions.Count(instruction => instruction.Operand is MethodReference method && method.Name == "SendEOI") != 2 ||
            dispatch.Body.Instructions.Any(instruction => instruction.Operand is MethodReference method && method.Name == "GetId"))
            throw new InvalidOperationException("The device-IRQ scheduler index or original EOI paths changed.");
        VerifyExport(hal, PhysicalMethod, PhysicalExport, true);
        VerifyExport(hal, SchedulerMethod, SchedulerExport, false);
    }

    public static string Fingerprint(MethodDefinition method)
    {
        string Operand(object? operand) => operand switch
        {
            Instruction branch => "#" + method.Body.Instructions.IndexOf(branch),
            VariableDefinition local => "$" + local.Index,
            MemberReference member => member.FullName,
            null => "",
            IFormattable value => value.ToString(null, CultureInfo.InvariantCulture),
            _ => operand.ToString()!
        };
        var value = string.Join("|", method.IsStatic, method.ReturnType.FullName,
            string.Join(";", method.Parameters.Select(parameter => parameter.ParameterType.FullName)),
            method.Body.InitLocals, method.Body.ExceptionHandlers.Count,
            string.Join(";", method.Body.Variables.Select(local => local.VariableType.FullName)),
            string.Join("\n", method.Body.Instructions.Select(instruction => instruction.OpCode.Code + ":" + Operand(instruction.Operand))));
        return Hash(Encoding.UTF8.GetBytes(value));
    }

    public static string Hash(byte[] bytes) => Convert.ToHexStringLower(SHA256.HashData(bytes));
    private static TypeDefinition Initializer(AssemblyDefinition hal) => hal.MainModule.Types.Single(type => type.FullName == "Cosmos.Kernel.HAL.X64.X64PlatformInitializer");
    public static MethodDefinition Count(AssemblyDefinition hal) => Initializer(hal).Methods.Single(method => method.Name == "GetCpuCount");
    public static MethodDefinition Timer(AssemblyDefinition core) => core.MainModule.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.X64.Cpu.LocalApic").Methods.Single(method => method.Name == "HandleTimerInterrupt");
    public static MethodDefinition Dispatch(AssemblyDefinition core) => core.MainModule.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.X64.Cpu.X64InterruptController").Methods.Single(method => method.Name == "Dispatch");
    private static void Versions(AssemblyDefinition core, AssemblyDefinition hal)
    {
        if (core.Name.Name != CoreName || hal.Name.Name != HalName || core.Name.Version != new Version(3, 0, 94, 0) || hal.Name.Version != new Version(3, 0, 94, 0))
            throw new InvalidOperationException("Only the pinned Cosmos x64 platform 3.0.94 is supported.");
        if (!Timer(core).HasBody || !Count(hal).HasBody || !Dispatch(core).HasBody || Timer(core).Parameters.Count != 1 ||
            Count(hal).Parameters.Count != 0 || Dispatch(core).Parameters.Count != 1)
            throw new InvalidOperationException("Pinned platform method signatures changed.");
    }

    private static void AddExport(AssemblyDefinition hal, string name, string entry, MethodReference? physical)
    {
        var module = hal.MainModule;
        var method = new MethodDefinition(name, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt32);
        var attribute = new CustomAttribute(module.ImportReference(typeof(UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
        attribute.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, entry)));
        method.CustomAttributes.Add(attribute);
        Initializer(hal).Methods.Add(method);
        var il = method.Body.GetILProcessor();
        il.Append(physical is null ? il.Create(OpCodes.Ldc_I4_1) : il.Create(OpCodes.Call, physical));
        il.Append(il.Create(OpCodes.Ret));
        method.Body.MaxStackSize = 1;
    }

    private static void VerifyExport(AssemblyDefinition hal, string name, string entry, bool physical)
    {
        var method = Initializer(hal).Methods.Single(candidate => candidate.Name == name);
        if (!method.IsAssembly || !method.IsStatic || method.ReturnType.MetadataType != MetadataType.UInt32 || method.Parameters.Count != 0 ||
            !method.HasBody || method.Body.InitLocals || method.Body.Variables.Count != 0 || method.Body.ExceptionHandlers.Count != 0 ||
            method.Body.Instructions.Count != 2 || method.Body.Instructions[1].OpCode != OpCodes.Ret ||
            method.CustomAttributes.Count != 1 || method.CustomAttributes[0].AttributeType.FullName != typeof(UnmanagedCallersOnlyAttribute).FullName ||
            method.CustomAttributes[0].Fields.Count != 1 || method.CustomAttributes[0].Fields[0].Name != "EntryPoint" ||
            !Equals(method.CustomAttributes[0].Fields[0].Argument.Value, entry))
            throw new InvalidOperationException("The platform diagnostic export signature changed.");
        var first = method.Body.Instructions[0];
        if (physical ? first.OpCode != OpCodes.Call || first.Operand is not MethodReference getter ||
                getter.FullName != "System.UInt32 Cosmos.Kernel.Core.X64.AcpiMadt::get_CpuCount()" : first.OpCode != OpCodes.Ldc_I4_1)
            throw new InvalidOperationException("The platform diagnostic is not its expected read-only leaf.");
    }
}