// Unit tests for analyzeGadgetEffect — the per-gadget symbolic executor. Builds
// synthetic raw ir::pcode for known x86-64 gadgets (no r2/Ghidra) and checks the
// classified transfer function.
#include <z3++.h>

#include <catch2/catch_test_macros.hpp>
#include <vector>

#include "fold/binary/gadget_semantics.hpp"
#include "fold/ir/pcode.hpp"

namespace {
using namespace fold::ir::pcode;
using namespace fold::binary;

// x86-64 SLEIGH register offsets: RAX=0x0, RDX=0x10, RSI=0x30, RDI=0x38,
// RSP=0x20; RIP modeled at 0x288. (Matches the Ghidra x86 register file.)
GadgetRegModel x64Model() {
    GadgetRegModel m;
    m.ptr_bits = 64;
    m.sp_offset = 0x20;
    m.args = {{"rdi", 0x38}, {"rsi", 0x30}, {"rdx", 0x10}};
    m.syscall_reg = "rax";
    m.syscall_off = 0x0;
    m.name_by_off = {{0x0, "rax"},  {0x10, "rdx"}, {0x20, "rsp"},
                     {0x30, "rsi"}, {0x38, "rdi"}, {0x288, "rip"}};
    return m;
}

Varnode reg(std::uint64_t off, std::uint32_t sz = 8) { return {Space::Register, off, sz}; }
Varnode con(std::uint64_t v, std::uint32_t sz = 8) { return {Space::Constant, v, sz}; }
Varnode sid() { return con(1, 8); } // LOAD/STORE space-id operand

PcodeOp mk(Op op, Varnode out, std::vector<Varnode> in) {
    PcodeOp p;
    p.op = op;
    p.out = out;
    p.in = std::move(in);
    return p;
}
PcodeOp mkNoOut(Op op, std::vector<Varnode> in) {
    PcodeOp p;
    p.op = op;
    p.in = std::move(in);
    return p;
}

// `pop rip ; ret`-style epilogue (rip = [rsp]; rsp += 8; return) appended to a
// gadget so it terminates like a real `ret`.
void appendRet(Function& fn) {
    fn.ops.push_back(mk(Op::Load, reg(0x288), {sid(), reg(0x20)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x20), {reg(0x20), con(8)}));
    fn.ops.push_back(mkNoOut(Op::Return, {reg(0x288)}));
}

GadgetEffect analyze(const Function& fn) {
    z3::context ctx;
    return analyzeGadgetEffect(fn, x64Model(), ctx);
}
} // namespace

TEST_CASE("gadget: pop rdi ; ret → rdi from stack[0], sp_delta 16", "[gadget]") {
    Function fn; // rdi = [rsp]; rsp += 8; <ret>
    fn.ops.push_back(mk(Op::Load, reg(0x38), {sid(), reg(0x20)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x20), {reg(0x20), con(8)}));
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.sets("rdi"));
    REQUIRE(e.reg_out["rdi"].src == ValSrc::StackSlot);
    REQUIRE(e.reg_out["rdi"].stack_offset == 0);
    REQUIRE(e.sp_delta == 16);
    REQUIRE(e.term == GadgetEffect::Term::Ret);
}

TEST_CASE("gadget: mov rax, 0x3b ; syscall → rax const, syscall terminator", "[gadget]") {
    Function fn;
    fn.ops.push_back(mk(Op::Copy, reg(0x0), {con(0x3b)}));
    fn.ops.push_back(mkNoOut(Op::CallOther, {con(0)}));
    auto e = analyze(fn);
    REQUIRE(e.reg_out["rax"].src == ValSrc::Const);
    REQUIRE(e.reg_out["rax"].constant == 0x3b);
    REQUIRE(e.term == GadgetEffect::Term::Syscall);
}

TEST_CASE("gadget: pop rsi ; pop rdx ; ret → both from stack, sp_delta 24", "[gadget]") {
    Function fn; // rsi=[rsp]; rsp+=8; rdx=[rsp]; rsp+=8; <ret>
    fn.ops.push_back(mk(Op::Load, reg(0x30), {sid(), reg(0x20)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x20), {reg(0x20), con(8)}));
    fn.ops.push_back(mk(Op::Load, reg(0x10), {sid(), reg(0x20)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x20), {reg(0x20), con(8)}));
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.reg_out["rsi"].src == ValSrc::StackSlot);
    REQUIRE(e.reg_out["rsi"].stack_offset == 0);
    REQUIRE(e.reg_out["rdx"].src == ValSrc::StackSlot);
    REQUIRE(e.reg_out["rdx"].stack_offset == 8);
    REQUIRE(e.sp_delta == 24);
}

TEST_CASE("gadget: mov [rdi], rsi ; ret → write-what-where (reg→reg)", "[gadget]") {
    Function fn;
    fn.ops.push_back(mkNoOut(Op::Store, {sid(), reg(0x38), reg(0x30)})); // [rdi] = rsi
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.mem_writes.size() == 1);
    REQUIRE(e.mem_writes[0].addr.src == ValSrc::Reg);
    REQUIRE(e.mem_writes[0].addr.reg == "rdi");
    REQUIRE(e.mem_writes[0].value.src == ValSrc::Reg);
    REQUIRE(e.mem_writes[0].value.reg == "rsi");
}

TEST_CASE("gadget: pop rdi ; add rdi,0x10 ; ret → StackSlot[0] with addend 0x10", "[gadget]") {
    Function fn; // rdi=[rsp]; rsp+=8; rdi += 0x10; <ret>
    fn.ops.push_back(mk(Op::Load, reg(0x38), {sid(), reg(0x20)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x20), {reg(0x20), con(8)}));
    fn.ops.push_back(mk(Op::IntAdd, reg(0x38), {reg(0x38), con(0x10)}));
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.reg_out["rdi"].src == ValSrc::StackSlot);
    REQUIRE(e.reg_out["rdi"].stack_offset == 0);
    REQUIRE(e.reg_out["rdi"].addend == 0x10);
}

TEST_CASE("gadget: mov rsi, rdi ; ret → reg-move (Reg source, addend 0)", "[gadget]") {
    Function fn; // rsi = rdi
    fn.ops.push_back(mk(Op::Copy, reg(0x30), {reg(0x38)}));
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.reg_out["rsi"].src == ValSrc::Reg);
    REQUIRE(e.reg_out["rsi"].reg == "rdi");
    REQUIRE(e.reg_out["rsi"].addend == 0);
}

TEST_CASE("gadget: add rdi,0x10 ; ret → self-additive (Reg rdi, addend 0x10, not skipped)",
          "[gadget]") {
    Function fn; // rdi += 0x10 (no stack pop) — a real change, must not be elided
    fn.ops.push_back(mk(Op::IntAdd, reg(0x38), {reg(0x38), con(0x10)}));
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE(e.sets("rdi"));
    REQUIRE(e.reg_out["rdi"].src == ValSrc::Reg);
    REQUIRE(e.reg_out["rdi"].reg == "rdi");
    REQUIRE(e.reg_out["rdi"].addend == 0x10);
}

TEST_CASE("gadget: bare ret → no arg regs changed, sp_delta 8", "[gadget]") {
    Function fn;
    appendRet(fn);
    auto e = analyze(fn);
    REQUIRE_FALSE(e.sets("rdi"));
    REQUIRE_FALSE(e.sets("rsi"));
    REQUIRE(e.sp_delta == 8);
    REQUIRE(e.term == GadgetEffect::Term::Ret);
}