using System;
using System.Runtime.InteropServices;
using WhiteMagic.Native;

namespace WhiteMagic.Hooking;

/// <summary>
/// A single reversible inline detour. Replaces the start of a native function
/// with a jump to a managed hook delegate, preserves the overwritten bytes in a
/// remote trampoline, and exposes the trampoline through <see cref="CallOriginal"/>.
/// </summary>
/// <remarks>
/// Only supported in-process. The detour uses a 5-byte relative <c>jmp</c> on x86
/// targets and a 14-byte RIP-relative absolute <c>jmp</c> on x64 targets.
/// </remarks>
public sealed class Detour : IDisposable
{
    private readonly MemoryBase _memory;

    /// <summary>The unique name of this detour.</summary>
    public string Name { get; }

    /// <summary>The target native function address.</summary>
    public IntPtr Target { get; }

    /// <summary>The managed hook delegate that the detour invokes.</summary>
    public Delegate Hook { get; }

    /// <summary>The bytes overwritten at <see cref="Target"/>.</summary>
    public byte[] OverwrittenBytes { get; private set; } = Array.Empty<byte>();

    /// <summary>
    /// The allocated trampoline that executes the original prologue and then jumps
    /// back into the original function.
    /// </summary>
    public IntPtr Trampoline { get; private set; }

    /// <summary>
    /// A delegate wrapping <see cref="Trampoline"/> with the same type as <see cref="Hook"/>.
    /// </summary>
    public Delegate? Original { get; private set; }

    /// <summary><see langword="true"/> while the detour bytes are live at <see cref="Target"/>.</summary>
    public bool IsApplied { get; private set; }

    internal Detour(MemoryBase memory, string name, IntPtr target, Delegate hook)
    {
        ArgumentNullException.ThrowIfNull(hook);

        _memory = memory;
        Name = name;
        Target = target;
        Hook = hook;
    }

    /// <summary>
    /// Installs the detour after validating that the required overwrite covers whole
    /// prologue instructions.
    /// </summary>
    /// <exception cref="InvalidOperationException">The prologue cannot be safely spliced.</exception>
    public void Apply()
    {
        if (IsApplied)
            return;

        int pointerSize = _memory.Is64Bit ? 8 : 4;
        int detourLength = pointerSize == 8 ? 14 : 5;

        // Try to read detourLength + 16 bytes for the prologue decoder.
        // If the target is near a page boundary, this might fail, so fall back to the minimum.
        byte[] prologue = _memory.ReadBytes(Target, detourLength + 16);
        if (prologue.Length < detourLength)
        {
            // Second attempt: read only the minimum required bytes
            prologue = _memory.ReadBytes(Target, detourLength);
            if (prologue.Length < detourLength)
            {
                throw new InvalidOperationException(
                    "Could not read enough bytes from the target function to install a detour.");
            }
        }

        int preserveLength = PrologueDecoder.GetWholeInstructionLength(prologue, detourLength, _memory.Is64Bit);
        OverwrittenBytes = new byte[preserveLength];
        Buffer.BlockCopy(prologue, 0, OverwrittenBytes, 0, preserveLength);

        IntPtr hookAddress = Marshal.GetFunctionPointerForDelegate(Hook);
        byte[] hookJump = pointerSize == 8
            ? BuildAbsoluteJump(hookAddress)
            : BuildRelativeJump(Target, hookAddress);

        // Allocate and build the trampoline before touching the target.
        int returnJumpSize = pointerSize == 8 ? 14 : 5;
        int trampolineSize = preserveLength + returnJumpSize;
        IntPtr trampoline = NativeMethods.VirtualAllocEx(
            _memory.Handle,
            IntPtr.Zero,
            trampolineSize,
            MemoryAllocationType.Commit | MemoryAllocationType.Reserve,
            MemoryProtectionType.ExecuteReadWrite);

        if (trampoline == IntPtr.Zero)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new InvalidOperationException(
                $"Failed to allocate detour trampoline: error {error}");
        }

        try
        {
            var trampolineBytes = new byte[trampolineSize];
            OverwrittenBytes.CopyTo(trampolineBytes, 0);

            byte[] returnJump = pointerSize == 8
                ? BuildAbsoluteJump(Target + preserveLength)
                : BuildRelativeJump(trampoline + preserveLength, Target + preserveLength);

            returnJump.CopyTo(trampolineBytes, preserveLength);

            int written = _memory.WriteBytes(trampoline, trampolineBytes);
            if (written != trampolineSize)
            {
                throw new InvalidOperationException(
                    "Failed to write the detour trampoline into the target process.");
            }

            // Make the target page writable if necessary, then write the detour jump,
            // restoring the original protection regardless of success or failure.
            if (!NativeMethods.VirtualProtectEx(
                _memory.Handle,
                Target,
                preserveLength,
                MemoryProtectionType.ExecuteReadWrite,
                out MemoryProtectionType oldProtect))
            {
                int error = Marshal.GetLastPInvokeError();
                throw new InvalidOperationException(
                    $"Failed to change target memory protection: error {error}");
            }

            try
            {
                written = _memory.WriteBytes(Target, hookJump);
                if (written != hookJump.Length)
                    throw new InvalidOperationException("Failed to write detour jump to target.");

                Trampoline = trampoline;
                Original = Marshal.GetDelegateForFunctionPointer(Trampoline, Hook.GetType());
                IsApplied = true;
            }
            finally
            {
                NativeMethods.VirtualProtectEx(
                    _memory.Handle,
                    Target,
                    preserveLength,
                    oldProtect,
                    out _);
            }
        }
        catch
        {
            NativeMethods.VirtualFreeEx(
                _memory.Handle,
                trampoline,
                0,
                MemoryFreeType.Release);
            throw;
        }
    }

    /// <summary>Restores the original bytes and releases the trampoline.</summary>
    public void Remove()
    {
        if (!IsApplied)
            return;

        if (OverwrittenBytes.Length > 0 && Target != IntPtr.Zero)
        {
            NativeMethods.VirtualProtectEx(
                _memory.Handle,
                Target,
                OverwrittenBytes.Length,
                MemoryProtectionType.ExecuteReadWrite,
                out MemoryProtectionType oldProtect);

            try
            {
                _memory.WriteBytes(Target, OverwrittenBytes);
            }
            finally
            {
                NativeMethods.VirtualProtectEx(
                    _memory.Handle,
                    Target,
                    OverwrittenBytes.Length,
                    oldProtect,
                    out _);
            }
        }

        if (Trampoline != IntPtr.Zero)
        {
            NativeMethods.VirtualFreeEx(
                _memory.Handle,
                Trampoline,
                0,
                MemoryFreeType.Release);
        }

        Trampoline = IntPtr.Zero;
        Original = null;
        OverwrittenBytes = Array.Empty<byte>();
        IsApplied = false;
    }

    /// <summary>
    /// Invokes the original function through the trampoline. Pass the same arguments
    /// that the native signature expects; the return value is boxed.
    /// </summary>
    public object? CallOriginal(params object?[] args)
    {
        if (Original is null)
        {
            throw new InvalidOperationException(
                "The detour is not applied; there is no original trampoline to call.");
        }

        return Original.DynamicInvoke(args);
    }

    /// <inheritdoc />
    public void Dispose()
    {
        Remove();
    }

    private static byte[] BuildRelativeJump(IntPtr source, IntPtr destination)
    {
        byte[] bytes = new byte[5];
        bytes[0] = 0xE9;

        long distance = (long)destination - ((long)source + 5);
        if (distance < int.MinValue || distance > int.MaxValue)
        {
            throw new ArgumentOutOfRangeException(nameof(destination),
                "Relative jump distance exceeds the 2 GiB range of an E8/E9 encoding.");
        }

        uint rel = (uint)distance;
        bytes[1] = (byte)rel;
        bytes[2] = (byte)(rel >> 8);
        bytes[3] = (byte)(rel >> 16);
        bytes[4] = (byte)(rel >> 24);
        return bytes;
    }

    private static byte[] BuildAbsoluteJump(IntPtr destination)
    {
        // jmp [rip+0] followed by the absolute target address.
        byte[] bytes = new byte[14];
        bytes[0] = 0xFF;
        bytes[1] = 0x25;
        bytes[2] = 0x00;
        bytes[3] = 0x00;
        bytes[4] = 0x00;
        bytes[5] = 0x00;

        long addr = (long)destination;
        bytes[6] = (byte)addr;
        bytes[7] = (byte)(addr >> 8);
        bytes[8] = (byte)(addr >> 16);
        bytes[9] = (byte)(addr >> 24);
        bytes[10] = (byte)(addr >> 32);
        bytes[11] = (byte)(addr >> 40);
        bytes[12] = (byte)(addr >> 48);
        bytes[13] = (byte)(addr >> 56);
        return bytes;
    }
}