using System.Diagnostics;
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
    ];

    // 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];

    [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_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);

        // The executor will allocate a remote call stub; our reader then refuses every
        // WriteBytes call. The allocation made before the failure must still be freed.
        // The write failure must surface as an InvalidOperationException, not hang or crash.
        Assert.Throws<InvalidOperationException>(() =>
            executor.Execute<int>(new IntPtr(0x123456789ABCDEF0L), CallConvention.Stdcall));
    }

    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.
    /// </summary>
    private sealed class WriteFailingMemoryBase : 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)
            => 0;

        public override void Dispose()
        {
        }
    }
}