using System;
using System.Diagnostics;
using WhiteMagic;
using WhiteMagic.Assembly;
using WhiteMagic.Native;
using Xunit;

namespace WhiteMagicTest;

/// <summary>
/// Tests for <see cref="RemoteModule"/> / <see cref="RemoteFunction"/> resolution and
/// execution through the <see cref="Magic"/> facade (task 7.2).
/// </summary>
public class ModuleFunctionTests
{
    // Ensure a module is loaded in this process before resolving it.
    private static IntPtr Load(string module)
    {
        IntPtr handle = NativeMethods.LoadLibrary(module);
        Assert.NotEqual(IntPtr.Zero, handle);
        return handle;
    }

    [Fact]
    public void Module_indexer_resolves_base_address()
    {
        IntPtr handle = Load("kernel32.dll");

        using var magic = Magic.OpenInProcess();
        RemoteModule module = magic["kernel32"];

        // The module handle returned by LoadLibrary is the module's base address.
        Assert.Equal(handle, module.BaseAddress);
        Assert.Equal("KERNEL32.DLL", module.Name, ignoreCase: true);
    }

    [Fact]
    public void Function_indexer_resolves_direct_export()
    {
        IntPtr handle = Load("user32.dll");
        IntPtr expected = NativeMethods.GetProcAddress(handle, "MessageBoxA");
        Assert.NotEqual(IntPtr.Zero, expected);

        using var magic = Magic.OpenInProcess();
        RemoteFunction fn = magic["user32"]["MessageBoxA"];

        Assert.Equal(expected, fn.Address);
        Assert.Equal("MessageBoxA", fn.Name);
    }

    [Fact]
    public void Function_indexer_follows_export_forwarder()
    {
        // kernel32!HeapAlloc is a classic forwarder to NTDLL.RtlAllocateHeap. Whatever the
        // OS loader resolves it to, our parser must reach the same final address.
        IntPtr handle = Load("kernel32.dll");
        IntPtr expected = NativeMethods.GetProcAddress(handle, "HeapAlloc");
        Assert.NotEqual(IntPtr.Zero, expected);

        using var magic = Magic.OpenInProcess();
        RemoteFunction fn = magic["kernel32"]["HeapAlloc"];

        Assert.Equal(expected, fn.Address);
    }

    [Fact]
    public void Module_indexer_throws_for_unloaded_module()
    {
        using var magic = Magic.OpenInProcess();
        Assert.Throws<DllNotFoundException>(() => magic["definitely-not-loaded-xyz.dll"]);
    }

    [Fact]
    public void Function_indexer_throws_for_unknown_export()
    {
        Load("kernel32.dll");

        using var magic = Magic.OpenInProcess();
        Assert.Throws<InvalidOperationException>(() => magic["kernel32"]["NoSuchExport_ZZZ"]);
    }

    private delegate uint GetCurrentProcessIdDelegate();

    [Fact]
    public void CreateDelegate_throws_for_external_session()
    {
        Load("kernel32.dll");

        // External reader (even to self): the address is not treated as host-mapped, so a
        // delegate to it is rejected rather than handed back to AV on invocation.
        using var magic = Magic.Open(Process.GetCurrentProcess());
        RemoteFunction fn = magic["kernel32"]["GetCurrentProcessId"];

        Assert.Throws<InvalidOperationException>(() => fn.CreateDelegate<GetCurrentProcessIdDelegate>());
    }

    [Fact]
    public void CreateDelegate_invokes_function_in_process()
    {
        Load("kernel32.dll");

        using var magic = Magic.OpenInProcess();
        var getPid = magic["kernel32"]["GetCurrentProcessId"].CreateDelegate<GetCurrentProcessIdDelegate>();

        Assert.Equal((uint)Process.GetCurrentProcess().Id, getPid());
    }

    [Fact]
    public void Resolved_function_executes_via_remote_thread()
    {
        Load("kernel32.dll");

        using var magic = Magic.OpenInProcess();
        RemoteFunction getPid = magic["kernel32"]["GetCurrentProcessId"];

        // GetCurrentProcessId takes no args and is thread-agnostic; a remote thread in our
        // own process must report our PID.
        uint pid = getPid.Execute<uint>(CallConvention.Stdcall);
        Assert.Equal((uint)Process.GetCurrentProcess().Id, pid);
    }
}