namespace WhiteMagic.Assembly;

/// <summary>
/// The default <see cref="IAssembler"/> backend. Hand-emits calling-convention
/// trampolines and remote-execution stubs using deterministic byte emitters
/// (<see cref="EmitU8"/>, <see cref="EmitU32"/>, <see cref="EmitU64"/>). Has
/// no native or third-party dependency — no FASM, no Iced.
/// </summary>
/// <remarks>
/// <see cref="Assemble"/> is not supported by this backend (it is a parse-free
/// emitter, not a text assembler). Use <see cref="IcedAssembler"/> (Phase 8) for
/// arbitrary mnemonics.
/// </remarks>
public sealed class StubAssembler : IAssembler
{
    /// <summary>
    /// The maximum number of call arguments a single stub may pass. Far above any real
    /// calling convention; exists only to keep frame-size and stack-offset math from
    /// overflowing <see cref="int"/>.
    /// </summary>
    public const int MaxArguments = 256;

    /// <inheritdoc />
    public byte[] Assemble(string assemblyText, ulong origin = 0)
    {
        throw new NotSupportedException(
            "StubAssembler does not parse text assembly. " +
            "Use IcedAssembler (Phase 8) for arbitrary mnemonics.");
    }

    // ── Emit primitives ────────────────────────────────────────────────────

    public void EmitU8(List<byte> buffer, byte value) => buffer.Add(value);

    public void EmitU32(List<byte> buffer, uint value)
    {
        buffer.Add((byte)value);
        buffer.Add((byte)(value >> 8));
        buffer.Add((byte)(value >> 16));
        buffer.Add((byte)(value >> 24));
    }

    public void EmitU64(List<byte> buffer, ulong value)
    {
        EmitU32(buffer, (uint)value);
        EmitU32(buffer, (uint)(value >> 32));
    }

    // ── Call-stub builders ─────────────────────────────────────────────────

    /// <summary>
    /// Builds a calling-convention call stub for x86 or x64.
    /// </summary>
    /// <param name="stubAddress">Where the stub lands (for E8 rel32 encoding).</param>
    /// <param name="targetAddress">Function to call.</param>
    /// <param name="arguments">Argument values. For x86 each element holds a 32-bit argument;
    /// for x64 each element holds the full 64-bit pointer-sized argument.</param>
    /// <param name="pointerSize">4 (x86) or 8 (x64).</param>
    /// <param name="convention">Calling convention (ignored on x64; Windows has a single ABI).</param>
    /// <exception cref="ArgumentOutOfRangeException"><paramref name="pointerSize"/> is not 4 or 8,
    /// or <paramref name="convention"/> is not known, or <paramref name="arguments"/> exceeds
    /// <see cref="MaxArguments"/>, or the distance between stub and target exceeds the E8 rel32
    /// range.</exception>
    public byte[] BuildCallStub(IntPtr stubAddress, IntPtr targetAddress,
        nuint[] arguments, int pointerSize, CallConvention convention)
    {
        ArgumentNullException.ThrowIfNull(arguments);
        // Bound the argument count. No real calling convention passes anywhere near this
        // many; the cap keeps the frame-size and stack-offset arithmetic (0x20 + 8*stackArgs)
        // well inside int range, so it can never overflow into a bogus/negative frame.
        ArgumentOutOfRangeException.ThrowIfGreaterThan(arguments.Length, MaxArguments, nameof(arguments));

        var buffer = new List<byte>(96);

        if (pointerSize == 4)
        {
            // X86 args are 32-bit. Truncate nuint down to uint — callers must pass values
            // that fit in 32 bits on x86 targets.
            uint[] args32 = new uint[arguments.Length];
            for (int i = 0; i < arguments.Length; i++)
            {
                ulong v = arguments[i];
                if (v > uint.MaxValue)
                {
                    throw new ArgumentOutOfRangeException(nameof(arguments),
                        $"Argument {i} = 0x{v:X} does not fit in 32 bits (x86 target).");
                }
                args32[i] = (uint)v;
            }
            BuildX86Stub(buffer, checked((uint)stubAddress), checked((uint)targetAddress),
                args32, convention);
        }
        else if (pointerSize == 8)
        {
            // Windows x64 uses a single ABI — the convention parameter is unused.
            BuildX64Stub(buffer, (ulong)(nint)stubAddress, (ulong)(nint)targetAddress, arguments);
        }
        else
        {
            throw new ArgumentOutOfRangeException(nameof(pointerSize), pointerSize,
                $"Expected 4 (x86) or 8 (x64), got {pointerSize}.");
        }

        return buffer.ToArray();
    }

    private void BuildX86Stub(List<byte> buffer, uint stubAddr,
        uint target, uint[] args, CallConvention convention)
    {
        uint current = stubAddr;
        int argIndex = 0;

        switch (convention)
        {
            case CallConvention.Thiscall when args.Length - argIndex >= 1:
                EmitMovRegImm32(buffer, 0xB9, args[argIndex], ref current); // mov ecx, arg0
                argIndex++;
                break;

            case CallConvention.Fastcall:
                if (args.Length - argIndex >= 1)
                {
                    EmitMovRegImm32(buffer, 0xB9, args[argIndex], ref current); // mov ecx, arg0
                    argIndex++;
                }
                if (args.Length - argIndex >= 1)
                {
                    EmitMovRegImm32(buffer, 0xBA, args[argIndex], ref current); // mov edx, arg1
                    argIndex++;
                }
                break;

            case CallConvention.Cdecl:
            case CallConvention.Stdcall:
                break;

            default:
                throw new ArgumentOutOfRangeException(nameof(convention), convention,
                    $"Unsupported calling convention: {convention}.");
        }

        // Push remaining args in reverse order (right-to-left)
        for (int i = args.Length - 1; i >= argIndex; i--)
        {
            current += 5;
            buffer.Add(0x68);             // push imm32
            EmitU32(buffer, args[i]);
        }

        // call rel32
        long distance = (long)target - (long)(current + 5);
        if (distance < int.MinValue || distance > int.MaxValue)
        {
            throw new ArgumentOutOfRangeException(
                $"target (0x{target:X}) is >2 GiB from stub (0x{stubAddr:X}); " +
                "E8 rel32 cannot encode this distance. Place the stub closer to the target.");
        }
        buffer.Add(0xE8);
        EmitU32(buffer, (uint)distance);
        current += 5;

        // Caller cleanup (cdecl only)
        int stackCount = args.Length - argIndex;
        if (convention == CallConvention.Cdecl && stackCount > 0)
        {
            int cleanup = stackCount * 4;
            if (cleanup <= 127)
            {
                buffer.Add(0x83);          // add esp, imm8
                buffer.Add(0xC4);
                buffer.Add((byte)cleanup);
            }
            else
            {
                buffer.Add(0x81);          // add esp, imm32
                buffer.Add(0xC4);
                EmitU32(buffer, (uint)cleanup);
            }
        }

        buffer.Add(0xC3);                 // ret
    }

    /// <summary>
    /// Builds a Windows x64 call stub that conforms to the Microsoft x64 ABI:
    /// first 4 integer/pointer args in RCX, RDX, R8, R9 (64-bit loads); stack args
    /// above a 32-byte shadow space; 16-byte stack alignment at the inner <c>call</c>.
    /// </summary>
    /// <remarks>
    /// <para>Frame derivation. The ABI requires the <em>inner</em> <c>call</c> site to
    /// land with call-site rsp ≡ 0 (mod 16), so that <c>call</c> pushes 8 bytes and the
    /// callee sees entry rsp ≡ 8 — the value an MSVC prologue (<c>push rbp; sub rsp, 0x20</c>)
    /// expects, and the only value for which locals land 16-aligned (SSE-safe).</para>
    /// <list type="bullet">
    /// <item>Stub entry: rsp ≡ 8 (mod 16).</item>
    /// <item>Need post-sub rsp ≡ 0 → sub operand K satisfies K ≡ 8 (mod 16).</item>
    /// <item>Frame must hold shadow space (0x20) + stack args (8 bytes each for args 4+).
    /// Choose the smallest such K: <c>K = frameBytes + ((8 − frameBytes) mod 16 + 16) mod 16</c>.
    /// For 0–5 args, K ∈ {0x28, 0x38}; pattern scales linearly.</item>
    /// </list>
    /// <code>
    /// sub rsp, K                    ; K ≡ 8 (mod 16), K ≥ 0x20 + 8·stackArgs
    /// mov rcx, arg0     ; REX.W + imm64    (10 bytes)
    /// mov rdx, arg1     ; REX.W + imm64    (10 bytes)
    /// mov r8,  arg2     ; REX.WB+ imm64    (10 bytes, REX.R)
    /// mov r9,  arg3     ; REX.WB+ imm64    (10 bytes, REX.R)
    /// mov rax, arg[N]   ; REX.W + imm64    (10 bytes)
    /// mov [rsp + 0x20 + 8*(N-4)], rax      (5/8 bytes)
    /// call target (rel32)                  ( 5 bytes)
    /// add rsp, K                           ( 7 bytes)
    /// ret                                  ( 1 byte)
    /// </code>
    /// </remarks>
    private void BuildX64Stub(List<byte> buffer, ulong stubAddr,
        ulong target, nuint[] args)
    {
        ulong current = stubAddr;

        // Compute frame size K. K ≡ 8 (mod 16) so that the inner call sees
        // post-sub rsp ≡ 0 and delivers target entry rsp ≡ 8 (mod 16).
        int stackArgs = Math.Max(0, args.Length - 4);
        int frameBytes = 0x20 + 8 * stackArgs;
        int k = frameBytes + ((8 - (frameBytes % 16) + 16) % 16);

        // sub rsp, imm32 (always imm32 form — constant 7 bytes regardless of K).
        buffer.Add(0x48); buffer.Add(0x81); buffer.Add(0xEC);
        EmitU32(buffer, (uint)k);
        current += 7;

        // 64-bit register loads for args 0..3. All encodings are exactly 10 bytes:
        //   REX.W (0x48) + 0xB9 + imm64    → mov rcx, imm64
        //   REX.W (0x48) + 0xBA + imm64    → mov rdx, imm64
        //   REX.WB(0x49) + 0xB8 + imm64    → mov r8,  imm64 (REX.R for r8)
        //   REX.WB(0x49) + 0xB9 + imm64    → mov r9,  imm64 (REX.R)
        byte[][] regMoves =
        [
            [0x48, 0xB9],
            [0x48, 0xBA],
            [0x49, 0xB8],
            [0x49, 0xB9],
        ];

        int regCount = Math.Min(args.Length, 4);
        for (int i = 0; i < regCount; i++)
        {
            byte[] prefix = regMoves[i];
            buffer.Add(prefix[0]);
            buffer.Add(prefix[1]);
            EmitU64(buffer, args[i]);
            current += (uint)(prefix.Length + 8);
        }

        // Stack args: written at [post-sub-rsp + 0x20 + 8*(i-4)], i.e. above the
        // shadow window, where the inner call's callee expects them.
        for (int i = 4; i < args.Length; i++)
        {
            int offset = 0x20 + (i - 4) * 8;
            buffer.Add(0x48); buffer.Add(0xB8);       // mov rax, imm64
            EmitU64(buffer, args[i]);
            current += 10;

            buffer.Add(0x48); buffer.Add(0x89);       // mov [rsp + disp], rax
            if (offset <= 127)
            {
                buffer.Add(0x44); buffer.Add(0x24);   // ModRM: [rsp + disp8]
                buffer.Add((byte)offset);
                current += 5;
            }
            else
            {
                buffer.Add(0x84); buffer.Add(0x24);   // ModRM: [rsp + disp32]
                EmitU32(buffer, (uint)offset);
                current += 8;
            }
        }

        // call rel32
        long distance = (long)target - (long)(current + 5);
        if (distance is < int.MinValue or > int.MaxValue)
        {
            throw new ArgumentOutOfRangeException(
                "target and stub are >2 GiB apart; E8 rel32 cannot encode this distance.");
        }
        buffer.Add(0xE8);
        EmitU32(buffer, (uint)distance);
        current += 5;

        // Tear down the frame symmetrically.
        buffer.Add(0x48); buffer.Add(0x81); buffer.Add(0xC4);
        EmitU32(buffer, (uint)k);

        buffer.Add(0xC3);
    }

    // ── Instruction helpers ────────────────────────────────────────────────

    /// <summary>Emit <c>mov reg32, imm32</c> and advances <paramref name="ip"/> by 5.</summary>
    private static void EmitMovRegImm32(List<byte> buffer, byte opcode, uint imm32, ref uint ip)
    {
        buffer.Add(opcode);
        buffer.Add((byte)imm32);
        buffer.Add((byte)(imm32 >> 8));
        buffer.Add((byte)(imm32 >> 16));
        buffer.Add((byte)(imm32 >> 24));
        ip += 5;
    }
}