XFESystem

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

/// <summary>
/// Explicit, allocation-free access to the existing firmware framebuffer writer.
/// The kernel keeps its serial trace separate and disables this writer once its
/// first desktop frame owns the framebuffer. No firmware services are invoked.
/// </summary>
internal static class EarlyBootConsoleRuntimePatch
{
    public const string SetEnabledName = "XfeBootConsoleSetEnabledV1";
    public const string WriteCharName = "XfeBootConsoleWriteCharV1";
    public const string SetEnabledExport = "XfeBootConsoleSetEnabled";
    public const string WriteCharExport = "XfeBootConsoleWriteChar";

    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var console = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.EarlyGop");
        var enabled = console.Fields.Single(field => field.Name == "Enabled");
        var putChar = console.Methods.Single(method => method.Name == "PutChar" && method.Parameters.Count == 1);
        if (!console.IsAbstract || !console.IsSealed || !enabled.IsStatic || !enabled.IsPublic || enabled.IsInitOnly ||
            enabled.FieldType.MetadataType != MetadataType.Boolean || !putChar.IsPublic || !putChar.IsStatic ||
            putChar.ReturnType.MetadataType != MetadataType.Void || putChar.Parameters[0].ParameterType.MetadataType != MetadataType.Char ||
            !putChar.HasBody || !putChar.Body.Instructions.Any(instruction => instruction.OpCode == OpCodes.Ldsfld &&
                instruction.Operand is FieldReference field && field.FullName == enabled.FullName))
            throw new InvalidOperationException("The pinned early framebuffer console contract changed.");

        var existing = console.Methods.Where(method => method.Name is SetEnabledName or WriteCharName).ToArray();
        if (existing.Length != 0)
        {
            if (existing.Length != 2) throw new InvalidOperationException("Incomplete early framebuffer console exports.");
            Verify();
            return true;
        }
        foreach (var type in module.Types)
            if (type.Methods.Any(method => method.CustomAttributes.Any(attribute => attribute.AttributeType.FullName ==
                    "System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute" && attribute.Fields.Any(field =>
                        field.Name == "EntryPoint" && field.Argument.Value is string entry && entry is SetEnabledExport or WriteCharExport))))
                throw new InvalidOperationException("An early framebuffer console export name is already in use.");

        var setter = Export(SetEnabledName, SetEnabledExport);
        var si = setter.Body.GetILProcessor();
        si.Append(si.Create(OpCodes.Ldarg_0));
        si.Append(si.Create(OpCodes.Ldc_I4_0));
        si.Append(si.Create(OpCodes.Cgt_Un));
        si.Append(si.Create(OpCodes.Stsfld, enabled));
        si.Append(si.Create(OpCodes.Ret));
        setter.Body.MaxStackSize = 2;

        var writer = Export(WriteCharName, WriteCharExport);
        var wi = writer.Body.GetILProcessor();
        var pass = wi.Create(OpCodes.Ldarg_0);
        var replace = wi.Create(OpCodes.Ldc_I4, (int)'?');
        var convert = wi.Create(OpCodes.Conv_U2);
        wi.Append(wi.Create(OpCodes.Ldarg_0));
        wi.Append(wi.Create(OpCodes.Ldc_I4, (int)'\n'));
        wi.Append(wi.Create(OpCodes.Beq, pass));
        wi.Append(wi.Create(OpCodes.Ldarg_0));
        wi.Append(wi.Create(OpCodes.Ldc_I4, (int)'\r'));
        wi.Append(wi.Create(OpCodes.Beq, pass));
        wi.Append(wi.Create(OpCodes.Ldarg_0));
        wi.Append(wi.Create(OpCodes.Ldc_I4, 0x20));
        wi.Append(wi.Create(OpCodes.Blt_Un, replace));
        wi.Append(wi.Create(OpCodes.Ldarg_0));
        wi.Append(wi.Create(OpCodes.Ldc_I4, 0x7e));
        wi.Append(wi.Create(OpCodes.Bgt_Un, replace));
        wi.Append(pass);
        wi.Append(wi.Create(OpCodes.Br, convert));
        wi.Append(replace);
        wi.Append(convert);
        wi.Append(wi.Create(OpCodes.Call, putChar));
        wi.Append(wi.Create(OpCodes.Ret));
        writer.Body.MaxStackSize = 2;
        Verify();
        return false;

        MethodDefinition Export(string name, string entryPoint)
        {
            var method = new MethodDefinition(name, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig,
                module.TypeSystem.Void);
            method.Parameters.Add(new("value", ParameterAttributes.None, module.TypeSystem.Int32));
            var attribute = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute)
                .GetConstructor(Type.EmptyTypes)!));
            attribute.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, entryPoint)));
            method.CustomAttributes.Add(attribute);
            console.Methods.Add(method);
            return method;
        }

        void Verify()
        {
            var setter = console.Methods.Single(method => method.Name == SetEnabledName);
            var writer = console.Methods.Single(method => method.Name == WriteCharName);
            foreach (var (method, entryPoint) in new[] { (setter, SetEnabledExport), (writer, WriteCharExport) })
                if (!method.IsStatic || !method.IsAssembly || method.ReturnType.MetadataType != MetadataType.Void ||
                    method.Parameters.Count != 1 || method.Parameters[0].ParameterType.MetadataType != MetadataType.Int32 ||
                    !method.HasBody || method.Body.InitLocals || method.Body.Variables.Count != 0 || method.Body.ExceptionHandlers.Count != 0 ||
                    method.CustomAttributes.Count != 1 || method.CustomAttributes[0].AttributeType.FullName !=
                        "System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute" || method.CustomAttributes[0].Fields.Count != 1 ||
                    method.CustomAttributes[0].Fields[0].Name != "EntryPoint" || !Equals(method.CustomAttributes[0].Fields[0].Argument.Value, entryPoint))
                    throw new InvalidOperationException("An early framebuffer console export contract changed.");
            var set = setter.Body.Instructions;
            if (set.Count != 5 || set[0].OpCode != OpCodes.Ldarg_0 || set[1].OpCode != OpCodes.Ldc_I4_0 ||
                set[2].OpCode != OpCodes.Cgt_Un || set[3].OpCode != OpCodes.Stsfld ||
                set[3].Operand is not FieldReference field || field.FullName != enabled.FullName || set[4].OpCode != OpCodes.Ret)
                throw new InvalidOperationException("The early framebuffer console setter is not its expected leaf.");
            var write = writer.Body.Instructions;
            var expected = new[] { OpCodes.Ldarg_0, OpCodes.Ldc_I4, OpCodes.Beq, OpCodes.Ldarg_0, OpCodes.Ldc_I4, OpCodes.Beq,
                OpCodes.Ldarg_0, OpCodes.Ldc_I4, OpCodes.Blt_Un, OpCodes.Ldarg_0, OpCodes.Ldc_I4, OpCodes.Bgt_Un,
                OpCodes.Ldarg_0, OpCodes.Br, OpCodes.Ldc_I4, OpCodes.Conv_U2, OpCodes.Call, OpCodes.Ret };
            if (write.Count != expected.Length || write.Where((instruction, index) => instruction.OpCode != expected[index]).Any() ||
                !Equals(write[1].Operand, (int)'\n') || !Equals(write[4].Operand, (int)'\r') || !Equals(write[7].Operand, 0x20) ||
                !Equals(write[10].Operand, 0x7e) || !Equals(write[14].Operand, (int)'?') ||
                !ReferenceEquals(write[2].Operand, write[12]) || !ReferenceEquals(write[5].Operand, write[12]) ||
                !ReferenceEquals(write[8].Operand, write[14]) || !ReferenceEquals(write[11].Operand, write[14]) ||
                !ReferenceEquals(write[13].Operand, write[15]) || write[16].Operand is not MethodReference call || call.FullName != putChar.FullName)
                throw new InvalidOperationException("The early framebuffer console character export is not its bounded ASCII leaf.");
        }
    }
}