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

namespace WhiteMagic.Injection;

/// <summary>
/// Injects raw machine code into a process's memory.
/// </summary>
public static class CodeInjector
{
    /// <summary>
    /// Injects code at a specific address.
    /// </summary>
    /// <param name="memory">The memory accessor.</param>
    /// <param name="address">The target address.</param>
    /// <param name="code">The machine code bytes to write.</param>
    /// <returns>The address the code was written to (same as <paramref name="address"/>).</returns>
    /// <exception cref="ArgumentException"><paramref name="code"/> is empty.</exception>
    /// <exception cref="Win32Exception">Write fails.</exception>
    public static IntPtr InjectAtAddress(MemoryBase memory, IntPtr address, byte[] code)
    {
        ArgumentNullException.ThrowIfNull(memory);
        ArgumentNullException.ThrowIfNull(code);

        if (code.Length == 0)
            throw new ArgumentException("Code cannot be empty.", nameof(code));

        if (address == IntPtr.Zero)
            throw new ArgumentException("Address cannot be zero.", nameof(address));

        // Write the code to the target address
        int written = memory.WriteBytes(address, code);
        if (written != code.Length)
        {
            int error = Marshal.GetLastPInvokeError();
            throw new Win32Exception(error,
                $"WriteProcessMemory failed at {address} (wrote {written} of {code.Length} bytes).");
        }

        return address;
    }

    /// <summary>
    /// Allocates executable memory and injects code into it.
    /// </summary>
    /// <param name="memory">The memory accessor.</param>
    /// <param name="code">The machine code bytes to inject.</param>
    /// <param name="protection">
    /// The memory protection. Defaults to <see cref="MemoryProtectionType.ExecuteReadWrite"/>.
    /// </param>
    /// <returns>
    /// The base address of the allocated memory containing the code.
    /// The caller is responsible for freeing this memory (e.g., via <see cref="AllocatedMemory.Dispose"/>).
    /// </returns>
    /// <exception cref="ArgumentException"><paramref name="code"/> is empty.</exception>
    /// <exception cref="Win32Exception">Allocation or write fails.</exception>
    public static AllocatedMemory Inject(
        MemoryBase memory,
        byte[] code,
        MemoryProtectionType protection = MemoryProtectionType.ExecuteReadWrite)
    {
        ArgumentNullException.ThrowIfNull(memory);
        ArgumentNullException.ThrowIfNull(code);

        if (code.Length == 0)
            throw new ArgumentException("Code cannot be empty.", nameof(code));

        // Allocate memory with the specified protection
        var allocated = new AllocatedMemory(memory, code.Length, protection);

        // Write the code to the allocated memory
        int written = memory.WriteBytes(allocated.BaseAddress, code);
        if (written != code.Length)
        {
            int error = Marshal.GetLastPInvokeError();
            allocated.Dispose();
            throw new Win32Exception(error,
                $"WriteProcessMemory failed (wrote {written} of {code.Length} bytes).");
        }

        return allocated;
    }
}