using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using WhiteMagic;
using WhiteMagic.Assembly;
using WhiteMagic.Execution;
using WhiteMagic.Native;

namespace WhiteMagicTest.Execution;

public sealed class RemoteThreadExecutorTests
{
    // x64 payloads. Live execution tests run only on x64 because the payloads use the
    // Microsoft x64 ABI (integer args in RCX, RDX, R8, R9, then stack at [rsp+0x28]).

    // mov eax, ecx
    // add eax, edx
    // ret
    private static readonly byte[] AddPayload = [0x89, 0xC8, 0x01, 0xD0, 0xC3];

    // mov eax, ecx
    // add eax, edx
    // add eax, r8d
    // add eax, r9d
    // add eax, [rsp+0x28]
    // ret
    private static readonly byte[] SumFivePayload =
    [
        0x89, 0xC8,
        0x01, 0xD0,
        0x44, 0x01, 0xC0,
        0x44, 0x01, 0xC8,
        0x03, 0x84, 0x24, 0x28, 0x00, 0x00, 0x00,
        0xC3
    ];

    // Five-arg callee that also executes an alignment-sensitive SSE instruction, proving
    // the stub delivers a 16-byte-aligned stack the CPU actually accepts (movaps #GPs on a
    // misaligned address) alongside correct register+stack argument placement.
    // sub rsp, 24           ; entry rsp ≡ 8 (mod 16) -> rsp ≡ 0 (16-aligned), giving a
    //                       ; 16-byte aligned scratch at [rsp..rsp+16) below the saved
    //                       ; return address ([rsp+24]) so the store leaves it intact
    // movaps [rsp], xmm0    ; aligned 16-byte store — faults unless rsp is 16-aligned
    // add rsp, 24           ; restore
    // mov eax, ecx
    // add eax, edx
    // add eax, r8d
    // add eax, r9d
    // add eax, [rsp+0x28]   ; 5th arg above the shadow space
    // ret
    private static readonly byte[] SseAlignedSumPayload =
    [
        0x48, 0x83, 0xEC, 0x18,
        0x0F, 0x29, 0x04, 0x24,
        0x48, 0x83, 0xC4, 0x18,
        0x89, 0xC8,
        0x01, 0xD0,
        0x44, 0x01, 0xC0,
        0x44, 0x01, 0xC8,
        0x03, 0x84, 0x24, 0x28, 0x00, 0x00, 0x00,
        0xC3
    ];

    // xor eax, eax
    // cmp byte ptr [rcx+rax], 0
    // je done
    // inc eax
    // jmp loop
    // done: ret
    private static readonly byte[] Utf8LengthPayload =
    [
        0x31, 0xC0,
        0x80, 0x3C, 0x01, 0x00,
        0x74, 0x04,
        0xFF, 0xC0,
        0xEB, 0xF6,
        0xC3
    ];

    // mov eax, [rcx]
    // add eax, [rcx+4]
    // ret
    private static readonly byte[] PointSumPayload = [0x8B, 0x01, 0x03, 0x41, 0x04, 0xC3];

    // Measures the callee's entry stack alignment without faulting.
    // mov eax, esp
    // add eax, 8
    // and eax, 0x0F
    // ret
    // Returns (rsp + 8) & 15, which is 0 iff callee entry rsp ≡ 8 (mod 16) — the
    // Microsoft x64 ABI guarantee the stub must deliver.
    private static readonly byte[] AlignProbePayload = [0x89, 0xE0, 0x83, 0xC0, 0x08, 0x83, 0xE0, 0x0F, 0xC3];

    [StructLayout(LayoutKind.Sequential)]
    private struct Point
    {
        public int X;
        public int Y;
    }

    [Fact]
    public void Execute_adds_two_integers()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        int result = RunPayload(AddPayload, CallConvention.Cdecl, 10, 32);
        Assert.Equal(42, result);
    }

    [Fact]
    public void Execute_sums_register_and_stack_arguments()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        int result = RunPayload(SumFivePayload, CallConvention.Cdecl, 1, 2, 3, 4, 5);
        Assert.Equal(15, result);
    }

    [Fact]
    public void Execute_delivers_16byte_aligned_stack_to_callee()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        // Stub entry rsp ≡ 8 → sub rsp, K (K ≡ 8) → rsp ≡ 0 → call → callee entry rsp ≡ 8.
        // So (rsp + 8) & 15 == 0 when the frame math is right; a bad K (e.g. 0x20) yields 8.
        int misalign = RunPayload(AlignProbePayload, CallConvention.Cdecl);
        Assert.Equal(0, misalign);
    }

    [Fact]
    public void Execute_runs_sse_callee_with_five_args()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        // Correct result (150) requires BOTH the 5th arg reaching [rsp+0x28] AND the
        // aligned movaps not faulting. A broken frame size/alignment either mis-sums or
        // #GPs in the callee.
        int result = RunPayload(SseAlignedSumPayload, CallConvention.Cdecl, 10, 20, 30, 40, 50);
        Assert.Equal(150, result);
    }

    [Fact]
    public void Execute_marshals_string_as_utf8_pointer()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        int result = RunPayload(Utf8LengthPayload, CallConvention.Cdecl, "hello");
        Assert.Equal(5, result);
    }

    [Fact]
    public void Execute_marshals_struct_as_pointer()
    {
        if (!Environment.Is64BitProcess)
        {
            return;
        }

        int result = RunPayload(PointSumPayload, CallConvention.Cdecl, new Point { X = 30, Y = 12 });
        Assert.Equal(42, result);
    }

    [Fact]
    public void Execute_throws_when_handle_is_invalid()
    {
        var executor = new RemoteThreadExecutor(new InvalidProcessReader());

        InvalidOperationException ex = Assert.Throws<InvalidOperationException>(() =>
        {
            executor.Execute<int>((IntPtr)0x1234, CallConvention.Cdecl);
        });

        Assert.Contains("handle", ex.Message, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public void Execute_throws_when_address_is_zero()
    {
        using var reader = new InProcessReader();
        var executor = new RemoteThreadExecutor(reader);

        ArgumentException ex = Assert.Throws<ArgumentException>(() =>
        {
            executor.Execute<int>(IntPtr.Zero, CallConvention.Cdecl);
        });

        Assert.Equal("address", ex.ParamName);
    }

    [Fact]
    public void InProcessReader_reports_current_process_bitness()
    {
        using var reader = new InProcessReader();
        Assert.Equal(Environment.Is64BitProcess, reader.Is64Bit);
    }

    [Fact]
    public void ExternalReader_reports_current_process_bitness()
    {
        using var reader = new ExternalReader(Process.GetCurrentProcess());
        Assert.Equal(Environment.Is64BitProcess, reader.Is64Bit);
    }

    [Fact]
    public void Execute_releases_allocated_remote_memory_on_write_failure()
    {
        using var reader = new WriteFailingMemoryBase();
        var executor = new RemoteThreadExecutor(reader);

        var allocated = new List<IntPtr>();
        var freed = new List<IntPtr>();
        nint next = 0x4000_0000;

        executor.RemoteAllocator = size => { var p = (IntPtr)(next += 0x1000); allocated.Add(p); return p; };
        executor.RemoteReleaser = p => freed.Add(p);

        // The string arg is marshalled to remote scratch FIRST, then its write fails.
        // Pre-fix the scratch was tracked only AFTER the write, so it escaped the finally
        // free and leaked. Post-fix every allocation is released on the failure path.
        Assert.Throws<InvalidOperationException>(() =>
            executor.Execute<int>(new IntPtr(0x123456789ABCDEF0L), CallConvention.Stdcall, "leakme"));

        Assert.NotEmpty(allocated);                                       // the arg scratch was allocated
        Assert.Equal(allocated.OrderBy(x => x), freed.OrderBy(x => x));   // and every alloc freed
    }

    private static int RunPayload(byte[] payload, CallConvention convention, params object?[] args)
    {
        const nint pageSize = 4096;
        const nint blockSize = pageSize * 2;

        using var reader = new InProcessReader();
        var executor = new RemoteThreadExecutor(reader);

        // Allocate a single executable block. The payload lives at the start and the
        // generated call stub is written to the second page, guaranteeing that the
        // relative CALL instruction stays within its ±2 GiB range.
        IntPtr block = NativeMethods.VirtualAllocEx(
            reader.Handle,
            IntPtr.Zero,
            blockSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);

        Assert.NotEqual(IntPtr.Zero, block);

        IntPtr stubAddress = block + pageSize;
        executor.StubAllocator = (_, size) => size <= pageSize ? stubAddress : IntPtr.Zero;

        try
        {
            int written = reader.WriteBytes(block, payload);
            Assert.Equal(payload.Length, written);

            return executor.Execute<int>(block, convention, args);
        }
        finally
        {
            NativeMethods.VirtualFreeEx(reader.Handle, block, 0, MemoryFreeType.Release);
        }
    }

    private sealed class InvalidProcessReader : MemoryBase
    {
        public override IntPtr ImageBase => IntPtr.Zero;

        public override SafeMemoryHandle Handle { get; } = new SafeMemoryHandle(new IntPtr(-1));

        public override bool Is64Bit => Environment.Is64BitProcess;

        public override int ProcessId => Environment.ProcessId;

        public override byte[] ReadBytes(IntPtr address, int count, bool isRelative = false)
            => throw new NotSupportedException();

        public override int WriteBytes(IntPtr address, ReadOnlySpan<byte> bytes, bool isRelative = false)
            => throw new NotSupportedException();

        public override void Dispose()
        {
        }
    }

    /// <summary>
    /// A fake reader whose WriteBytes always returns zero, forcing the executor down
    /// the failure path after it has allocated remote memory.
    /// Holds a valid handle to the current process so the executor passes its
    /// handle-validity check without performing real memory operations.
    /// </summary>
    private sealed class WriteFailingMemoryBase : MemoryBase
    {
        public override IntPtr ImageBase => IntPtr.Zero;
        public override SafeMemoryHandle Handle { get; }
        public override bool Is64Bit => Environment.Is64BitProcess;
        public override int ProcessId => Environment.ProcessId;

        public WriteFailingMemoryBase()
        {
            Handle = NativeMethods.OpenProcess(ProcessAccess.AllAccess, false, Environment.ProcessId);
        }

        public override byte[] ReadBytes(IntPtr address, int count, bool isRelative = false)
            => throw new NotSupportedException();

        public override int WriteBytes(IntPtr address, ReadOnlySpan<byte> bytes, bool isRelative = false)
            => 0;

        public override void Dispose()
        {
            Handle?.Dispose();
        }
    }
}