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

/// <summary>Serial diagnostics must never prevent a machine without COM1 from booting.</summary>
internal static class SerialRuntimePatch
{
    public const string HelperName = "XfeBoundedComWriteV1";
    // An iteration bound deliberately works before any clock, scheduler or IRQ exists.
    // A failed UART is disabled for this boot, so later diagnostics do not repeat it.
    public const int PollLimit = 4096;
    private const string Prefix = "XfeSerialV1_";
    private static readonly string[] CounterNames = ["Absent", "Timeout", "Dropped", "Written"];

    public static bool Apply(AssemblyDefinition assembly)
    {
        if (assembly.Name.Name != "Cosmos.Kernel.Core" || assembly.Name.Version != new Version(3, 0, 94, 0))
            throw new InvalidOperationException("Bounded serial output supports only pinned Cosmos.Kernel.Core 3.0.94.");
        var module = assembly.MainModule;
        var serial = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.Serial");
        var write = serial.Methods.Single(method => method.Name == "ComWrite" && method.IsStatic && method.HasBody &&
            method.ReturnType.MetadataType == MetadataType.Void && method.Parameters.Count == 1 && method.Parameters[0].ParameterType.MetadataType == MetadataType.Byte);
        var uart = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.CosmosFeatures").Methods.Single(method => method.Name == "get_UARTEnabled");
        var io = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Native").NestedTypes.Single(type => type.Name == "IO");
        var read = io.Methods.Single(method => method.Name == "Read8");
        var output = io.Methods.Single(method => method.Name == "Write8");
        var existing = serial.Methods.SingleOrDefault(method => method.Name == HelperName);
        if (existing is not null)
        {
            Verify(serial, write, uart, read, output);
            return true;
        }
        if (serial.Fields.Any(field => field.Name.StartsWith(Prefix, StringComparison.Ordinal)) ||
            serial.Methods.Any(method => method.Name.StartsWith("XfeSerial", StringComparison.Ordinal)))
            throw new InvalidOperationException("Incomplete or unknown bounded serial patch exists.");
        // Exact pinned x64 body, including branch destinations, constants and native calls.
        if (!OriginalReleaseShape(write, uart, read, output))
            throw new InvalidOperationException("Pinned x64 ComWrite IL changed; refusing to patch an unknown UART implementation.");
        var disabled = new FieldDefinition(Prefix + "Disabled", FieldAttributes.Private | FieldAttributes.Static, module.TypeSystem.Boolean);
        serial.Fields.Add(disabled);
        var counters = CounterNames.ToDictionary(name => name, name => new FieldDefinition(Prefix + name, FieldAttributes.Private | FieldAttributes.Static, module.TypeSystem.UInt64));
        foreach (var counter in counters.Values) serial.Fields.Add(counter);
        var helper = new MethodDefinition(HelperName, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Void);
        helper.Parameters.Add(new("value", ParameterAttributes.None, module.TypeSystem.Byte)); serial.Methods.Add(helper);
        EmitHelper(helper, disabled, counters, read, output);
        EmitWrapper(write, uart, helper);
        foreach (var name in CounterNames)
        {
            var getter = new MethodDefinition("XfeSerial" + name + "CountV1", MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt64);
            var attribute = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
            attribute.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, "XfeRuntimeSerial" + name + "Count")));
            getter.CustomAttributes.Add(attribute); serial.Methods.Add(getter);
            var il = getter.Body.GetILProcessor(); il.Append(il.Create(OpCodes.Ldsfld, counters[name])); il.Append(il.Create(OpCodes.Ret)); getter.Body.MaxStackSize = 1;
        }
        Verify(serial, write, uart, read, output);
        return false;
    }

    private static bool OriginalReleaseShape(MethodDefinition write, MethodReference uart, MethodReference read, MethodReference output)
    {
        var i = write.Body.Instructions;
        return write.Body.CodeSize == 34 && write.Body.Variables.Count == 0 && write.Body.ExceptionHandlers.Count == 0 && i.Count == 11 &&
            Calls(i[0], uart) && i[1].OpCode == OpCodes.Brfalse_S && ReferenceEquals(i[1].Operand, i[10]) &&
            i[2].OpCode == OpCodes.Ldc_I4 && Equals(i[2].Operand, 0x3fd) && Calls(i[3], read) &&
            i[4].OpCode == OpCodes.Ldc_I4_S && Equals(i[4].Operand, (sbyte)0x20) && i[5].OpCode == OpCodes.And &&
            i[6].OpCode == OpCodes.Brfalse_S && ReferenceEquals(i[6].Operand, i[2]) &&
            i[7].OpCode == OpCodes.Ldc_I4 && Equals(i[7].Operand, 0x3f8) && i[8].OpCode == OpCodes.Ldarg_0 &&
            Calls(i[9], output) && i[10].OpCode == OpCodes.Ret;
    }

    private static void EmitWrapper(MethodDefinition write, MethodReference uart, MethodReference helper)
    {
        write.Body = new MethodBody(write); var il = write.Body.GetILProcessor(); var end = il.Create(OpCodes.Ret);
        il.Append(il.Create(OpCodes.Call, uart)); il.Append(il.Create(OpCodes.Brfalse, end));
        il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Call, helper)); il.Append(end); write.Body.MaxStackSize = 1;
    }

    private static void EmitHelper(MethodDefinition helper, FieldReference disabled, IReadOnlyDictionary<string, FieldDefinition> counters, MethodReference read, MethodReference output)
    {
        helper.Body = new MethodBody(helper) { InitLocals = true, MaxStackSize = 3 };
        var count = new VariableDefinition(helper.Module.TypeSystem.Int32); var status = new VariableDefinition(helper.Module.TypeSystem.Byte);
        helper.Body.Variables.Add(count); helper.Body.Variables.Add(status);
        var il = helper.Body.GetILProcessor();
        var dropped = il.Create(OpCodes.Ldsfld, counters["Dropped"]); var absent = il.Create(OpCodes.Ldsfld, counters["Absent"]);
        var timeout = il.Create(OpCodes.Ldsfld, counters["Timeout"]); var disable = il.Create(OpCodes.Ldc_I4_1);
        var next = il.Create(OpCodes.Ldloc, count); var poll = il.Create(OpCodes.Ldc_I4, 0x3fd);
        il.Append(il.Create(OpCodes.Ldsfld, disabled)); il.Append(il.Create(OpCodes.Brtrue, dropped));
        il.Append(il.Create(OpCodes.Ldc_I4_0)); il.Append(il.Create(OpCodes.Stloc, count));
        il.Append(poll); il.Append(il.Create(OpCodes.Call, read)); il.Append(il.Create(OpCodes.Stloc, status));
        il.Append(il.Create(OpCodes.Ldloc, status)); il.Append(il.Create(OpCodes.Ldc_I4, 0xff)); il.Append(il.Create(OpCodes.Beq, absent));
        il.Append(il.Create(OpCodes.Ldloc, status)); il.Append(il.Create(OpCodes.Ldc_I4, 0x20)); il.Append(il.Create(OpCodes.And)); il.Append(il.Create(OpCodes.Brfalse, next));
        il.Append(il.Create(OpCodes.Ldc_I4, 0x3f8)); il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Call, output));
        Increment("Written"); il.Append(il.Create(OpCodes.Ret));
        il.Append(next); il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stloc, count));
        il.Append(il.Create(OpCodes.Ldloc, count)); il.Append(il.Create(OpCodes.Ldc_I4, PollLimit)); il.Append(il.Create(OpCodes.Blt, poll));
        il.Append(timeout); IncrementTail("Timeout"); il.Append(il.Create(OpCodes.Br, disable));
        il.Append(absent); IncrementTail("Absent");
        il.Append(disable); il.Append(il.Create(OpCodes.Stsfld, disabled));
        il.Append(dropped); IncrementTail("Dropped"); il.Append(il.Create(OpCodes.Ret));
        void Increment(string name) { il.Append(il.Create(OpCodes.Ldsfld, counters[name])); IncrementTail(name); }
        void IncrementTail(string name) { il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stsfld, counters[name])); }
    }

    private static void Verify(TypeDefinition serial, MethodDefinition write, MethodReference uart, MethodReference read, MethodReference output)
    {
        var helper = serial.Methods.Single(method => method.Name == HelperName);
        var disabled = serial.Fields.Single(field => field.Name == Prefix + "Disabled");
        var counters = CounterNames.ToDictionary(name => name, name => serial.Fields.Single(field => field.Name == Prefix + name));
        if (!disabled.IsPrivate || !disabled.IsStatic || disabled.FieldType.MetadataType != MetadataType.Boolean ||
            counters.Values.Any(field => !field.IsPrivate || !field.IsStatic || field.FieldType.MetadataType != MetadataType.UInt64) ||
            helper.ReturnType.MetadataType != MetadataType.Void || helper.Parameters.Count != 1 || helper.Parameters[0].ParameterType.MetadataType != MetadataType.Byte)
            throw new InvalidOperationException("Unexpected bounded serial state or signature.");
        var expected = new MethodDefinition("Expected", MethodAttributes.Static, serial.Module.TypeSystem.Void);
        expected.Parameters.Add(new("value", ParameterAttributes.None, serial.Module.TypeSystem.Byte)); serial.Methods.Add(expected);
        try
        {
            EmitHelper(expected, disabled, counters, read, output);
            if (Canonical(helper) != Canonical(expected)) throw new InvalidOperationException("Bounded serial helper differs from its fixed polling contract.");
            EmitWrapper(expected, uart, helper);
            if (Canonical(write) != Canonical(expected)) throw new InvalidOperationException("Bounded serial wrapper changed its feature gate.");
        }
        finally { serial.Methods.Remove(expected); }
        foreach (string name in CounterNames)
        {
            var getter = serial.Methods.Single(method => method.Name == "XfeSerial" + name + "CountV1");
            if (getter.Body.Instructions.Count != 2 || getter.Body.Instructions[0].OpCode != OpCodes.Ldsfld ||
                getter.Body.Instructions[0].Operand is not FieldReference field || field.FullName != counters[name].FullName ||
                getter.Body.Instructions[1].OpCode != OpCodes.Ret || !getter.CustomAttributes.Any(attribute => attribute.AttributeType.Name == "UnmanagedCallersOnlyAttribute" &&
                    attribute.Fields.Any(argument => argument.Name == "EntryPoint" && Equals(argument.Argument.Value, "XfeRuntimeSerial" + name + "Count"))))
                throw new InvalidOperationException("Serial diagnostic counter is not a read-only leaf export.");
        }
    }

    private static bool Calls(Instruction i, MethodReference method) => i.OpCode == OpCodes.Call && i.Operand is MethodReference call && call.FullName == method.FullName;
    private static string Canonical(MethodDefinition method) =>
        string.Join(";", method.Body.Variables.Select(variable => variable.VariableType.FullName)) + "|" + method.Body.InitLocals + "|" + method.Body.ExceptionHandlers.Count + "|" +
        string.Join("\n", method.Body.Instructions.Select(instruction => instruction.OpCode.Code + ":" + (instruction.Operand switch
        {
            Instruction branch => "#" + method.Body.Instructions.IndexOf(branch), VariableDefinition local => "$" + local.Index,
            MemberReference member => member.FullName, null => "", var operand => operand.ToString()
        })));
}