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

if (args.Length != 2) throw new ArgumentException("NativeInputPatch <private Cosmos.Kernel.System.dll> <kernel IL assembly>");
var runtimePath = Path.GetFullPath(args[0]);
var kernelPath = Path.GetFullPath(args[1]);
string? kernelProject = Path.GetDirectoryName(kernelPath);
while (kernelProject is not null && !File.Exists(Path.Combine(kernelProject, "XFESystem.Kernel.csproj"))) kernelProject = Path.GetDirectoryName(kernelProject);
if (kernelProject is null || Path.GetFileName(runtimePath) != "Cosmos.Kernel.System.dll" ||
    !runtimePath.StartsWith(Path.Combine(kernelProject, "obj") + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase))
    throw new ArgumentException("Only the project's private obj runtime copy may be instrumented.");
for (string? candidate = runtimePath; candidate is not null && !candidate.Equals(kernelProject, StringComparison.OrdinalIgnoreCase); candidate = Path.GetDirectoryName(candidate))
    if ((File.GetAttributes(candidate) & FileAttributes.ReparsePoint) != 0)
        throw new ArgumentException("A private runtime copy cannot be reached through a symbolic link or junction.");
var before = Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(runtimePath)));
using var runtime = AssemblyDefinition.ReadAssembly(runtimePath, new ReaderParameters { InMemory = true });
using var kernel = AssemblyDefinition.ReadAssembly(kernelPath, new ReaderParameters { InMemory = true });
var target = runtime.MainModule.Types.Single(t => t.FullName == "Cosmos.Kernel.System.Input.KeyboardManager")
    .Methods.Single(m => m.Name == "HandleScanCode" && m.IsStatic && m.HasBody && m.ReturnType.MetadataType == MetadataType.Void &&
        m.Parameters.Count == 2 && m.Parameters[0].ParameterType.MetadataType == MetadataType.Byte && m.Parameters[1].ParameterType.MetadataType == MetadataType.Boolean);
var capture = kernel.MainModule.Types.Single(t => t.FullName == "XFESystem.Kernel.NativeKeyboardCapture")
    .Methods.Single(m => m.Name == "Observe" && m.IsPublic && m.IsStatic && m.ReturnType.MetadataType == MetadataType.Void &&
        m.Parameters.Count == 2 && m.Parameters[0].ParameterType.MetadataType == MetadataType.Byte && m.Parameters[1].ParameterType.MetadataType == MetadataType.Boolean);
var existingHook = target.Body.Instructions.Where(i => i.Operand is MethodReference method && method.FullName == capture.FullName).ToArray();
if (existingHook.Length != 0 && (existingHook.Length != 1 || target.Body.Instructions.Count < 3 ||
    target.Body.Instructions[0].OpCode != OpCodes.Ldarg_0 || target.Body.Instructions[1].OpCode != OpCodes.Ldarg_1 ||
    target.Body.Instructions[2] != existingHook[0] || existingHook[0].OpCode != OpCodes.Call))
    throw new InvalidOperationException("An unexpected keyboard hook already exists; original decoder was not changed.");
if (existingHook.Length == 0)
{
    var il = target.Body.GetILProcessor();
    var first = target.Body.Instructions[0];
    il.InsertBefore(first, il.Create(OpCodes.Ldarg_0));
    il.InsertBefore(first, il.Create(OpCodes.Ldarg_1));
    il.InsertBefore(first, il.Create(OpCodes.Call, runtime.MainModule.ImportReference(capture)));
    target.Body.MaxStackSize = Math.Max(2, target.Body.MaxStackSize);
}
// Pointer hooks are a separate tail section. The keyboard prefix above and the
// original pointer decoder still execute unchanged; normal returns publish the
// final state instead of polling only the last state after a slow render frame.
var pointerType = runtime.MainModule.Types.Single(t => t.FullName == "Cosmos.Kernel.System.Input.MouseManager");
var pointerCaptureType = kernel.MainModule.Types.Single(t => t.FullName == "XFESystem.Kernel.NativePointerCapture");
var pointerCapture = pointerCaptureType.Methods.Single(m => m.Name == "Observe" && m.IsPublic && m.IsStatic &&
    m.ReturnType.MetadataType == MetadataType.Void && m.Parameters.Select(p => p.ParameterType.MetadataType)
        .SequenceEqual(new[] { MetadataType.Int32, MetadataType.Int32, MetadataType.Boolean, MetadataType.Boolean, MetadataType.Boolean, MetadataType.Int32, MetadataType.Int64 }));
var pointerClock = pointerCaptureType.Methods.Single(m => m.Name == "ReadTimestamp" && m.IsPublic && m.IsStatic &&
    m.ReturnType.MetadataType == MetadataType.Int64 && m.Parameters.Count == 0);
var pointerProperties = new[] { "X", "Y", "LeftButton", "RightButton", "MiddleButton" }
    .Select(name => pointerType.Methods.Single(m => m.Name == "get_" + name && m.IsStatic && m.Parameters.Count == 0 &&
        m.ReturnType.MetadataType == (name is "X" or "Y" ? MetadataType.Int32 : MetadataType.Boolean))).ToArray();
var pointerMethods = new[] { "HandleRelative", "HandleAbsolute" }.Select(name => pointerType.Methods.Single(m => m.Name == name && m.IsStatic && m.HasBody &&
    m.ReturnType.MetadataType == MetadataType.Void && m.Parameters.Count == (name == "HandleRelative" ? 4 : 3) &&
    m.Parameters[0].ParameterType.MetadataType == MetadataType.Int32 && m.Parameters[1].ParameterType.MetadataType == MetadataType.Int32 &&
    m.Parameters[2].ParameterType.FullName.EndsWith(".PointerButtons", StringComparison.Ordinal) &&
    (name != "HandleRelative" || m.Parameters[3].ParameterType.MetadataType == MetadataType.Int32))).ToArray();
bool pointerChanged = false;
foreach (var method in pointerMethods)
{
    var returns = method.Body.Instructions.Where(i => i.OpCode == OpCodes.Ret).ToArray();
    if (returns.Length == 0) throw new InvalidOperationException("A pointer decoder has no normal return path.");
    var hooks = method.Body.Instructions.Where(i => i.Operand is MethodReference called && called.FullName == pointerCapture.FullName).ToArray();
    if (hooks.Length != 0)
    {
        if (hooks.Length != returns.Length) throw new InvalidOperationException("An incomplete pointer hook already exists.");
        foreach (var ret in returns)
        {
            int index = method.Body.Instructions.IndexOf(ret);
            if (index < 9 || method.Body.Instructions[index - 9].OpCode != OpCodes.Nop ||
                method.Body.Instructions[index - 1].OpCode != OpCodes.Call || !Calls(method.Body.Instructions[index - 1], pointerCapture) ||
                !Calls(method.Body.Instructions[index - 2], pointerClock) ||
                method.Body.Instructions[index - 3].OpCode != (method.Name == "HandleRelative" ? OpCodes.Ldarg_3 : OpCodes.Ldc_I4_0) ||
                pointerProperties.Where((getter, position) => !Calls(method.Body.Instructions[index - 8 + position], getter)).Any())
                throw new InvalidOperationException("An unexpected pointer hook already exists; decoder was not changed.");
        }
        continue;
    }
    foreach (var ret in returns)
    {
        // Keep the original ret object as an anchor so every branch/switch/EH
        // boundary pointing at it enters the hook instead of bypassing it.
        ret.OpCode = OpCodes.Nop; ret.Operand = null;
        var il = method.Body.GetILProcessor(); var anchor = ret;
        void Add(Instruction instruction) { il.InsertAfter(anchor, instruction); anchor = instruction; }
        foreach (var getter in pointerProperties) Add(il.Create(OpCodes.Call, runtime.MainModule.ImportReference(getter)));
        Add(il.Create(method.Name == "HandleRelative" ? OpCodes.Ldarg_3 : OpCodes.Ldc_I4_0));
        Add(il.Create(OpCodes.Call, runtime.MainModule.ImportReference(pointerClock)));
        Add(il.Create(OpCodes.Call, runtime.MainModule.ImportReference(pointerCapture)));
        Add(il.Create(OpCodes.Ret));
    }
    foreach (var instruction in method.Body.Instructions)
        instruction.OpCode = LongBranch(instruction.OpCode);
    method.Body.MaxStackSize = Math.Max(7, method.Body.MaxStackSize);
    pointerChanged = true;
}
if (existingHook.Length == 0 || pointerChanged)
{
    var temporary = runtimePath + ".XfeInput.tmp";
    runtime.Write(temporary);
    File.Move(temporary, runtimePath, true);
}
var receipt = new
{
    Kind = "ProjectLocalRawKeyboardHook", RuntimePath = runtimePath, KernelPath = kernelPath,
    BeforeSha256 = before,
    AfterSha256 = Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(runtimePath))),
    OriginalDecoderPreserved = true, GlobalPackageCacheModified = false,
    PointerDecoderPreserved = true, PointerHooks = pointerMethods.Select(m => new { Method = m.FullName, NormalReturnHooks = m.Body.Instructions.Count(i => Calls(i, pointerCapture)) }).ToArray(),
    PointerTimestamp = "NativePointerCapture.ReadTimestamp / Stopwatch.GetTimestamp", WheelMeaning = "Relative packet delta after successful upstream decoding; absolute packets have zero wheel"
};
File.WriteAllText(runtimePath + ".XfeInput.json", JsonSerializer.Serialize(receipt, new JsonSerializerOptions { WriteIndented = true }));
Console.WriteLine("XFE native keyboard make/break hook ready.");
Console.WriteLine("XFE native pointer ordered button/motion/wheel hooks ready.");

static bool Calls(Instruction instruction, MethodReference method) => instruction.OpCode == OpCodes.Call && instruction.Operand is MethodReference called && called.FullName == method.FullName;
static OpCode LongBranch(OpCode opcode) => opcode.Code switch
{
    Code.Br_S => OpCodes.Br, Code.Brfalse_S => OpCodes.Brfalse, Code.Brtrue_S => OpCodes.Brtrue,
    Code.Beq_S => OpCodes.Beq, Code.Bge_S => OpCodes.Bge, Code.Bge_Un_S => OpCodes.Bge_Un, Code.Bgt_S => OpCodes.Bgt, Code.Bgt_Un_S => OpCodes.Bgt_Un,
    Code.Ble_S => OpCodes.Ble, Code.Ble_Un_S => OpCodes.Ble_Un, Code.Blt_S => OpCodes.Blt, Code.Blt_Un_S => OpCodes.Blt_Un,
    Code.Bne_Un_S => OpCodes.Bne_Un, Code.Leave_S => OpCodes.Leave, _ => opcode
};