using System;

namespace WhiteMagic.Hooking;

/// <summary>
/// Minimal instruction-length decoder for common x86/x64 prologue shapes.
/// The set is intentionally small: any opcode outside the covered set is rejected
/// rather than guessed. Full arbitrary-prologue validation is provided by the
/// optional Iced backend (Phase 8).
/// </summary>
/// <remarks>
/// Covered shapes:
/// <list type="bullet">
/// <item><c>push reg</c>: 0x50-0x57 (1 byte), including REX-prefixed forms.</item>
/// <item><c>push ebp/rbp</c>: 0x55 (1 byte).</item>
/// <item><c>mov edi, edi</c>: 8B FF (2 bytes).</item>
/// <item><c>mov ebp/rbp, esp/rsp</c>: 8B EC / 48 8B EC (2/3 bytes).</item>
/// <item><c>sub esp/rsp, imm8</c>: 83 EC imm8 / 48 83 EC imm8 (3/4 bytes).</item>
/// <item><c>sub esp/rsp, imm32</c>: 81 EC imm32 / 48 81 EC imm32 (6/7 bytes).</item>
/// </list>
/// </remarks>
internal static class PrologueDecoder
{
    /// <summary>
    /// Returns the length of the first instruction in <paramref name="bytes"/>
    /// if it matches a covered shape; otherwise returns -1.
    /// </summary>
    public static int GetInstructionLength(ReadOnlySpan<byte> bytes, bool is64Bit)
    {
        if (bytes.Length == 0)
            return 0;

        int i = 0;
        if (is64Bit && bytes[i] >= 0x40 && bytes[i] <= 0x4F)
        {
            // REX prefix.
            i++;
            if (bytes.Length <= i)
                return -1;
        }

        byte op = bytes[i];

        // push reg (0x50-0x57), including rbp (0x55).
        if ((op & 0xF8) == 0x50)
            return i + 1;

        // mov r32/64, r/m32/64. Recognize only the specific forms listed above.
        if (op == 0x8B && bytes.Length > i + 1)
        {
            byte modrm = bytes[i + 1];
            if (modrm == 0xFF || modrm == 0xEC)
                return i + 2;
        }

        // sub esp/rsp, imm8 — register-direct ModRM 0xEC only.
        if (op == 0x83 && bytes.Length > i + 2 && bytes[i + 1] == 0xEC)
            return i + 3;

        // sub esp/rsp, imm32 — register-direct ModRM 0xEC only.
        if (op == 0x81 && bytes.Length > i + 5 && bytes[i + 1] == 0xEC)
            return i + 6;

        return -1;
    }

    /// <summary>
    /// Walks prologue instructions until at least <paramref name="requiredBytes"/>
    /// have been covered, returning the total length of whole instructions that must
    /// be preserved in the trampoline.
    /// </summary>
    /// <exception cref="InvalidOperationException">An opcode is outside the covered set.</exception>
    public static int GetWholeInstructionLength(byte[] prologue, int requiredBytes, bool is64Bit)
    {
        int total = 0;
        while (total < requiredBytes)
        {
            int len = GetInstructionLength(prologue.AsSpan(total), is64Bit);
            if (len <= 0)
            {
                throw new InvalidOperationException(
                    "The target prologue contains an instruction outside the covered opcode set. " +
                    "Install the optional Iced backend for full instruction-boundary validation.");
            }

            total += len;
        }

        return total;
    }
}