XFESystem

using System.Reflection;
using System.Security.Cryptography;
using System.Text.Json;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Xunit;

namespace XFESystem.Desktop.Platform.Tests;

public sealed class NativeInputPatchTests
{
    private sealed record Fixture(string Root, string Runtime, string Kernel);

    private static Fixture Create(bool wrongCaptureSignature = false, bool misplacedHook = false)
    {
        var root = Path.Combine(AppContext.BaseDirectory, "PatchFixtures", Guid.NewGuid().ToString("N"));
        Directory.CreateDirectory(Path.Combine(root, "obj", "Native"));
        Directory.CreateDirectory(Path.Combine(root, "bin"));
        File.WriteAllText(Path.Combine(root, "XFESystem.Kernel.csproj"), "<Project />");
        string runtimePath = Path.Combine(root, "obj", "Native", "Cosmos.Kernel.System.dll");
        string kernelPath = Path.Combine(root, "bin", "XFESystem.Kernel.dll");
        using var kernel = AssemblyDefinition.CreateAssembly(new("XFESystem.Kernel", new Version(1, 0)), "XFESystem.Kernel", ModuleKind.Dll);
        var captureType = new TypeDefinition("XFESystem.Kernel", "NativeKeyboardCapture", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Abstract | Mono.Cecil.TypeAttributes.Sealed, kernel.MainModule.TypeSystem.Object);
        kernel.MainModule.Types.Add(captureType);
        var capture = new MethodDefinition("Observe", Mono.Cecil.MethodAttributes.Public | Mono.Cecil.MethodAttributes.Static, kernel.MainModule.TypeSystem.Void);
        capture.Parameters.Add(new("scanCode", Mono.Cecil.ParameterAttributes.None, wrongCaptureSignature ? kernel.MainModule.TypeSystem.Int32 : kernel.MainModule.TypeSystem.Byte));
        capture.Parameters.Add(new("released", Mono.Cecil.ParameterAttributes.None, kernel.MainModule.TypeSystem.Boolean));
        capture.Body.Instructions.Add(Instruction.Create(OpCodes.Ret)); captureType.Methods.Add(capture);
        AddPointerCapture(kernel);
        kernel.Write(kernelPath);
        using var runtime = AssemblyDefinition.CreateAssembly(new("Cosmos.Kernel.System", new Version(3, 0, 94, 0)), "Cosmos.Kernel.System", ModuleKind.Dll);
        var keyboard = new TypeDefinition("Cosmos.Kernel.System.Input", "KeyboardManager", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Abstract | Mono.Cecil.TypeAttributes.Sealed, runtime.MainModule.TypeSystem.Object);
        runtime.MainModule.Types.Add(keyboard);
        var decoder = new MethodDefinition("HandleScanCode", Mono.Cecil.MethodAttributes.Assembly | Mono.Cecil.MethodAttributes.Static, runtime.MainModule.TypeSystem.Void);
        decoder.Parameters.Add(new("scanCode", Mono.Cecil.ParameterAttributes.None, runtime.MainModule.TypeSystem.Byte));
        decoder.Parameters.Add(new("released", Mono.Cecil.ParameterAttributes.None, runtime.MainModule.TypeSystem.Boolean));
        keyboard.Methods.Add(decoder);
        decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_7));
        decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Pop));
        if (misplacedHook)
        {
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_0));
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_1));
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Call, runtime.MainModule.ImportReference(capture)));
        }
        decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
        AddPointerDecoder(runtime);
        runtime.Write(runtimePath);
        return new(root, runtimePath, kernelPath);
    }

    private static void AddPointerCapture(AssemblyDefinition kernel)
    {
        var type = new TypeDefinition("XFESystem.Kernel", "NativePointerCapture", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Abstract | Mono.Cecil.TypeAttributes.Sealed, kernel.MainModule.TypeSystem.Object);
        kernel.MainModule.Types.Add(type);
        var capture = new MethodDefinition("Observe", Mono.Cecil.MethodAttributes.Public | Mono.Cecil.MethodAttributes.Static, kernel.MainModule.TypeSystem.Void);
        foreach (var parameter in new[] { kernel.MainModule.TypeSystem.Int32, kernel.MainModule.TypeSystem.Int32, kernel.MainModule.TypeSystem.Boolean, kernel.MainModule.TypeSystem.Boolean,
            kernel.MainModule.TypeSystem.Boolean, kernel.MainModule.TypeSystem.Int32, kernel.MainModule.TypeSystem.Int64 }) capture.Parameters.Add(new(parameter));
        capture.Body.Instructions.Add(Instruction.Create(OpCodes.Ret)); type.Methods.Add(capture);
        var clock = new MethodDefinition("ReadTimestamp", Mono.Cecil.MethodAttributes.Public | Mono.Cecil.MethodAttributes.Static, kernel.MainModule.TypeSystem.Int64);
        clock.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I8, 123L)); clock.Body.Instructions.Add(Instruction.Create(OpCodes.Ret)); type.Methods.Add(clock);
    }

    private static void AddPointerDecoder(AssemblyDefinition runtime)
    {
        var buttons = new TypeDefinition("Cosmos.Kernel.HAL.Devices.Input", "PointerButtons", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Sealed,
            runtime.MainModule.ImportReference(typeof(Enum)));
        buttons.Fields.Add(new("value__", Mono.Cecil.FieldAttributes.Public | Mono.Cecil.FieldAttributes.SpecialName | Mono.Cecil.FieldAttributes.RTSpecialName, runtime.MainModule.TypeSystem.Int32));
        runtime.MainModule.Types.Add(buttons);
        var mouse = new TypeDefinition("Cosmos.Kernel.System.Input", "MouseManager", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Abstract | Mono.Cecil.TypeAttributes.Sealed, runtime.MainModule.TypeSystem.Object);
        runtime.MainModule.Types.Add(mouse);
        foreach (string property in new[] { "X", "Y", "LeftButton", "RightButton", "MiddleButton" })
        {
            var getter = new MethodDefinition("get_" + property, Mono.Cecil.MethodAttributes.Public | Mono.Cecil.MethodAttributes.Static,
                property is "X" or "Y" ? runtime.MainModule.TypeSystem.Int32 : runtime.MainModule.TypeSystem.Boolean);
            getter.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_0)); getter.Body.Instructions.Add(Instruction.Create(OpCodes.Ret)); mouse.Methods.Add(getter);
        }
        foreach (string name in new[] { "HandleRelative", "HandleAbsolute" })
        {
            var decoder = new MethodDefinition(name, Mono.Cecil.MethodAttributes.Assembly | Mono.Cecil.MethodAttributes.Static, runtime.MainModule.TypeSystem.Void);
            decoder.Parameters.Add(new(runtime.MainModule.TypeSystem.Int32)); decoder.Parameters.Add(new(runtime.MainModule.TypeSystem.Int32)); decoder.Parameters.Add(new(buttons));
            if (name == "HandleRelative") decoder.Parameters.Add(new(runtime.MainModule.TypeSystem.Int32));
            var firstReturn = Instruction.Create(OpCodes.Ret); var secondReturn = Instruction.Create(OpCodes.Ret);
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_0)); decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Brfalse_S, firstReturn));
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_7)); decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Pop));
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Br_S, secondReturn)); decoder.Body.Instructions.Add(firstReturn);
            decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Ldc_I4_8)); decoder.Body.Instructions.Add(Instruction.Create(OpCodes.Pop)); decoder.Body.Instructions.Add(secondReturn);
            mouse.Methods.Add(decoder);
        }
    }

    private static void RunPatch(string runtime, string kernel)
    {
        var main = Assembly.Load("NativeInputPatch").EntryPoint!;
        try { main.Invoke(null, [new[] { runtime, kernel }]); }
        catch (TargetInvocationException error) when (error.InnerException is not null)
        { System.Runtime.ExceptionServices.ExceptionDispatchInfo.Capture(error.InnerException).Throw(); }
    }

    private static string Hash(string path) => Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(path)));

    [Fact]
    public void PrivateCopyGetsOnePrefixAndOriginalDecoderInstructionsRemain()
    {
        var fixture = Create(); string kernelHash = Hash(fixture.Kernel), original = Hash(fixture.Runtime);
        RunPatch(fixture.Runtime, fixture.Kernel);
        using var patched = AssemblyDefinition.ReadAssembly(fixture.Runtime);
        var decoder = patched.MainModule.Types.Single(t => t.Name == "KeyboardManager").Methods.Single();
        Assert.Equal([OpCodes.Ldarg_0, OpCodes.Ldarg_1, OpCodes.Call, OpCodes.Ldc_I4_7, OpCodes.Pop, OpCodes.Ret], decoder.Body.Instructions.Select(i => i.OpCode).ToArray());
        var call = Assert.IsAssignableFrom<MethodReference>(decoder.Body.Instructions[2].Operand);
        Assert.Equal("XFESystem.Kernel.NativeKeyboardCapture", call.DeclaringType.FullName);
        Assert.Equal("Observe", call.Name); Assert.Equal(kernelHash, Hash(fixture.Kernel));
        using var receipt = JsonDocument.Parse(File.ReadAllText(fixture.Runtime + ".XfeInput.json"));
        Assert.Equal(original, receipt.RootElement.GetProperty("BeforeSha256").GetString());
        Assert.True(receipt.RootElement.GetProperty("OriginalDecoderPreserved").GetBoolean());
        Assert.False(receipt.RootElement.GetProperty("GlobalPackageCacheModified").GetBoolean());
    }

    [Fact]
    public void RepeatedPatchingIsByteForByteIdempotent()
    {
        var fixture = Create(); RunPatch(fixture.Runtime, fixture.Kernel); string first = Hash(fixture.Runtime);
        RunPatch(fixture.Runtime, fixture.Kernel); Assert.Equal(first, Hash(fixture.Runtime));
    }

    [Fact]
    public void RuntimeOutsidePrivateObjIsRejectedWithoutMutation()
    {
        var fixture = Create(); string outside = Path.Combine(fixture.Root, "Cosmos.Kernel.System.dll");
        File.Copy(fixture.Runtime, outside); string original = Hash(outside);
        Assert.Throws<ArgumentException>(() => RunPatch(outside, fixture.Kernel));
        Assert.Equal(original, Hash(outside)); Assert.False(File.Exists(outside + ".XfeInput.json"));
    }

    [Fact]
    public void IncorrectObserverSignatureIsRejectedWithoutMutation()
    {
        var fixture = Create(wrongCaptureSignature: true); string original = Hash(fixture.Runtime);
        Assert.Throws<InvalidOperationException>(() => RunPatch(fixture.Runtime, fixture.Kernel));
        Assert.Equal(original, Hash(fixture.Runtime));
    }

    [Fact]
    public void UnexpectedExistingHookIsRejectedWithoutMutation()
    {
        var fixture = Create(misplacedHook: true); string original = Hash(fixture.Runtime);
        Assert.Throws<InvalidOperationException>(() => RunPatch(fixture.Runtime, fixture.Kernel));
        Assert.Equal(original, Hash(fixture.Runtime));
    }

    [Fact]
    public void EveryPointerNormalReturnObservesFinalPropertiesPacketWheelAndMonotonicClock()
    {
        var fixture = Create(); RunPatch(fixture.Runtime, fixture.Kernel);
        using var patched = AssemblyDefinition.ReadAssembly(fixture.Runtime);
        var mouse = patched.MainModule.Types.Single(t => t.Name == "MouseManager");
        foreach (var decoder in mouse.Methods.Where(m => m.Name.StartsWith("Handle", StringComparison.Ordinal)))
        {
            var returns = decoder.Body.Instructions.Where(i => i.OpCode == OpCodes.Ret).ToArray(); Assert.Equal(2, returns.Length);
            foreach (var ret in returns)
            {
                int index = decoder.Body.Instructions.IndexOf(ret); Assert.Equal(OpCodes.Nop, decoder.Body.Instructions[index - 9].OpCode);
                Assert.Equal(new[] { "get_X", "get_Y", "get_LeftButton", "get_RightButton", "get_MiddleButton" },
                    decoder.Body.Instructions.Skip(index - 8).Take(5).Select(i => ((MethodReference)i.Operand).Name));
                Assert.Equal(decoder.Name == "HandleRelative" ? OpCodes.Ldarg_3 : OpCodes.Ldc_I4_0, decoder.Body.Instructions[index - 3].OpCode);
                Assert.Equal("ReadTimestamp", ((MethodReference)decoder.Body.Instructions[index - 2].Operand).Name);
                Assert.Equal("XFESystem.Kernel.NativePointerCapture", ((MethodReference)decoder.Body.Instructions[index - 1].Operand).DeclaringType.FullName);
            }
            var branches = decoder.Body.Instructions.Where(i => i.OpCode == OpCodes.Brfalse || i.OpCode == OpCodes.Br).ToArray(); Assert.Equal(2, branches.Length);
            Assert.All(branches, branch => Assert.Equal(OpCodes.Nop, ((Instruction)branch.Operand).OpCode));
            Assert.Equal(2, decoder.Body.Instructions.Count(i => i.OpCode == OpCodes.Pop));
            Assert.Contains(decoder.Body.Instructions, i => i.OpCode == OpCodes.Ldc_I4_7); Assert.Contains(decoder.Body.Instructions, i => i.OpCode == OpCodes.Ldc_I4_8);
        }
        using var receipt = JsonDocument.Parse(File.ReadAllText(fixture.Runtime + ".XfeInput.json"));
        Assert.True(receipt.RootElement.GetProperty("PointerDecoderPreserved").GetBoolean());
        Assert.Equal(2, receipt.RootElement.GetProperty("PointerHooks").GetArrayLength());
    }

    [Fact]
    public void MalformedPointerCaptureRejectsBeforeChangingAnyDecoderBytes()
    {
        var fixture = Create();
        using (var kernel = AssemblyDefinition.ReadAssembly(fixture.Kernel, new ReaderParameters { InMemory = true }))
        {
            kernel.MainModule.Types.Single(t => t.Name == "NativePointerCapture").Methods.Single(m => m.Name == "Observe").Parameters.RemoveAt(6);
            kernel.Write(fixture.Kernel);
        }
        string original = Hash(fixture.Runtime); Assert.Throws<InvalidOperationException>(() => RunPatch(fixture.Runtime, fixture.Kernel));
        Assert.Equal(original, Hash(fixture.Runtime)); Assert.False(File.Exists(fixture.Runtime + ".XfeInput.json"));
    }

    [Fact]
    public void ActualCosmos394PrivateCopyRetainsDecodedInstructionsAndIsIdempotent()
    {
        var fixture = Create();
        string packages = Environment.GetEnvironmentVariable("NUGET_PACKAGES") ?? Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".nuget", "packages");
        string upstream = Path.Combine(packages, "cosmos.kernel.system", "3.0.94", "runtimes", "linux-x64", "lib", "net10.0", "Cosmos.Kernel.System.dll");
        string upstreamHash = Hash(upstream);
        File.Copy(upstream, fixture.Runtime, true);
        File.Copy(typeof(XFESystem.Kernel.NativePointerCapture).Assembly.Location, fixture.Kernel, true);
        using var original = AssemblyDefinition.ReadAssembly(fixture.Runtime, new ReaderParameters { InMemory = true });
        var expected = original.MainModule.Types.Single(t => t.Name == "MouseManager").Methods.Where(m => m.Name is "HandleRelative" or "HandleAbsolute")
            .ToDictionary(m => m.Name, m => NormalizeDecoder(m));
        var originalKeyboard = NormalizeDecoder(original.MainModule.Types.Single(t => t.Name == "KeyboardManager").Methods.Single(m => m.Name == "HandleScanCode"));
        RunPatch(fixture.Runtime, fixture.Kernel); string firstHash = Hash(fixture.Runtime);
        using (var patched = AssemblyDefinition.ReadAssembly(fixture.Runtime))
        {
            foreach (var decoder in patched.MainModule.Types.Single(t => t.Name == "MouseManager").Methods.Where(m => expected.ContainsKey(m.Name)))
                Assert.Equal(expected[decoder.Name], NormalizeDecoder(decoder));
            Assert.Equal(originalKeyboard, NormalizeDecoder(patched.MainModule.Types.Single(t => t.Name == "KeyboardManager").Methods.Single(m => m.Name == "HandleScanCode")));
        }
        RunPatch(fixture.Runtime, fixture.Kernel); Assert.Equal(firstHash, Hash(fixture.Runtime)); Assert.Equal(upstreamHash, Hash(upstream));
    }

    private static string[] NormalizeDecoder(MethodDefinition method)
    {
        var original = new List<(Instruction Instruction, OpCode Code)>(); var body = method.Body.Instructions;
        for (int index = 0; index < body.Count; index++)
        {
            if (index == 0 && body.Count >= 3 && body[2].Operand is MethodReference keyboard && keyboard.DeclaringType.FullName == "XFESystem.Kernel.NativeKeyboardCapture")
            { index = 2; continue; }
            if (body[index].OpCode == OpCodes.Nop && index + 9 < body.Count && body[index + 8].Operand is MethodReference pointer &&
                pointer.DeclaringType.FullName == "XFESystem.Kernel.NativePointerCapture" && pointer.Name == "Observe" && body[index + 9].OpCode == OpCodes.Ret)
            { original.Add((body[index], OpCodes.Ret)); index += 9; continue; }
            original.Add((body[index], body[index].OpCode));
        }
        var positions = original.Select((item, index) => (item.Instruction, index)).ToDictionary(item => item.Instruction, item => item.index);
        return original.Select(item =>
        {
            string opcode = item.Code.Name;
            if (item.Code.OperandType == OperandType.ShortInlineBrTarget) opcode = opcode[..^2];
            string operand = item.Code == OpCodes.Ret ? "" : item.Instruction.Operand switch
            {
                Instruction branch => "target:" + positions[branch], Instruction[] branches => "targets:" + string.Join(',', branches.Select(branch => positions[branch])),
                MethodReference called => called.FullName, FieldReference field => field.FullName, TypeReference type => type.FullName,
                ParameterDefinition parameter => "arg:" + parameter.Index, VariableDefinition variable => "local:" + variable.Index,
                null => "", var value => value.ToString() ?? ""
            };
            return opcode + "|" + operand;
        }).ToArray();
    }
}