using System.Linq;
using Iced.Intel;
using WhiteMagic;
using WhiteMagic.Assembly;
using WhiteMagic.Hooking;

namespace WhiteMagicTest.Assembly;

/// <summary>
/// Tests for the optional <see cref="IcedAssembler"/> backend (tasks 8.1–8.3): arbitrary
/// text assembly, origin-relative encoding, and full prologue instruction decoding.
/// </summary>
public class IcedAssemblerTests
{
    private static Instruction[] Disassemble(byte[] code, int bitness, ulong origin)
    {
        var decoder = Decoder.Create(bitness, new ByteArrayCodeReader(code));
        decoder.IP = origin;

        var result = new List<Instruction>();
        ulong end = origin + (ulong)code.Length;
        while (decoder.IP < end)
            result.Add(decoder.Decode());

        return result.ToArray();
    }

    [Fact]
    public void Assemble_emits_single_instruction()
    {
        var assembler = new IcedAssembler(64);
        byte[] code = assembler.Assemble("ret");
        Assert.Equal(new byte[] { 0xC3 }, code);
    }

    [Fact]
    public void Assemble_emits_multiple_instructions_with_operands()
    {
        var assembler = new IcedAssembler(32);

        // The scenario from the managed-assembler spec.
        byte[] code = assembler.Assemble("push 0\nadd esp, 4\nret");
        Assert.NotEmpty(code);

        Instruction[] instructions = Disassemble(code, 32, 0);
        Assert.Equal(3, instructions.Length);
        Assert.Equal(Mnemonic.Push, instructions[0].Mnemonic);
        Assert.Equal(Mnemonic.Add, instructions[1].Mnemonic);
        Assert.Equal(Register.ESP, instructions[1].Op0Register);
        Assert.Equal(4UL, instructions[1].GetImmediate(1));
        Assert.Equal(Mnemonic.Ret, instructions[2].Mnemonic);
    }

    [Fact]
    public void Assemble_supports_comments_and_blank_lines()
    {
        var assembler = new IcedAssembler(64);
        byte[] code = assembler.Assemble("  ; prologue\n\nnop ; a comment\nret\n");

        Instruction[] instructions = Disassemble(code, 64, 0);
        Assert.Equal(2, instructions.Length);
        Assert.Equal(Mnemonic.Nop, instructions[0].Mnemonic);
        Assert.Equal(Mnemonic.Ret, instructions[1].Mnemonic);
    }

    [Fact]
    public void Assemble_encodes_label_branch_relative_to_origin()
    {
        var assembler = new IcedAssembler(64);
        const ulong origin = 0x1_4000_1000UL;

        // jmp forward over a nop to a label; the near-branch target must be resolved
        // against the supplied origin, not zero.
        byte[] code = assembler.Assemble("jmp done\nnop\ndone:\nret", origin);

        Instruction[] instructions = Disassemble(code, 64, origin);
        Instruction jmp = instructions[0];
        Assert.Equal(Mnemonic.Jmp, jmp.Mnemonic);

        // Target = origin + len(jmp) + len(nop): the address of the 'done: ret'.
        ulong expected = origin + (ulong)jmp.Length + 1;
        Assert.Equal(expected, jmp.NearBranchTarget);
    }

    [Fact]
    public void Assemble_throws_on_unsupported_operand()
    {
        var assembler = new IcedAssembler(64);
        Assert.Throws<NotSupportedException>(() => assembler.Assemble("mov rax, [rbx]"));
    }

    [Fact]
    public void GetPrologueLength_decodes_prologue_the_builtin_decoder_rejects()
    {
        // 48 8B C1 = mov rax, rcx — a register-to-register mov the built-in PrologueDecoder
        // does not cover (it only recognizes the 8B FF / 8B EC forms).
        // Followed by push rbp; mov rbp,rsp; sub rsp,0x20; mov rax,rcx to exceed 14 bytes.
        byte[] prologue =
        [
            0x48, 0x8B, 0xC1,             // mov rax, rcx      (3)
            0x55,                         // push rbp          (1)
            0x48, 0x8B, 0xEC,             // mov rbp, rsp      (3)
            0x48, 0x83, 0xEC, 0x20,       // sub rsp, 0x20     (4)
            0x48, 0x8B, 0xC1              // mov rax, rcx      (3)  -> total 14
        ];

        // The built-in decoder refuses the very first instruction.
        Assert.Throws<InvalidOperationException>(() =>
            PrologueDecoder.GetWholeInstructionLength(prologue, 14, is64Bit: true));

        // The Iced backend decodes it and returns the whole-instruction length covering
        // at least the 14 bytes a detour needs.
        var iced = new IcedAssembler();
        int length = iced.GetPrologueLength(prologue, 14, is64Bit: true);
        Assert.Equal(14, length);
    }

    [Fact]
    public void DetourManager_prologue_resolver_defaults_to_builtin_and_is_replaceable()
    {
        using var reader = new InProcessReader();
        var manager = new DetourManager(reader);

        // Default resolver is the built-in decoder.
        Assert.Throws<InvalidOperationException>(() =>
            manager.PrologueLengthResolver(new byte[] { 0x48, 0x8B, 0xC1, 0x90, 0x90 }, 4, true));

        // Swapping in the Iced resolver validates the same bytes.
        manager.PrologueLengthResolver = new IcedAssembler().GetPrologueLength;
        int length = manager.PrologueLengthResolver(new byte[] { 0x48, 0x8B, 0xC1, 0x90, 0x90 }, 4, true);
        Assert.True(length >= 4);
    }
}