//! Integration tests for the activation-record driver (control flow + calls).
//!
//! Exercises the structured-control-flow instructions (`block`, `loop`,
//! `if`/`else`, `br`/`br_if`/`br_table`, `return`), direct and indirect calls,
//! host-function bridging, and the configurable stack-overflow trap.

use we_wasm::runtime::{
    Engine, FuncAddr, FunctionInstance, HostFunc, ModuleAddr, ModuleInstance, Ref, Store, Table,
    Trap, WasmValue,
};
use we_wasm::{BlockType, FuncType, Instruction, Limits, RefType, TableType, ValType};

// ---------------------------------------------------------------------------
// Harness — small helpers for building modules in tests.
// ---------------------------------------------------------------------------

struct ModuleBuilder {
    store: Store,
    module_addr: ModuleAddr,
}

impl ModuleBuilder {
    fn new() -> Self {
        let mut store = Store::new();
        let module_addr = store.alloc_module(ModuleInstance::new());
        Self { store, module_addr }
    }

    fn with_types(types: Vec<FuncType>) -> Self {
        let mut b = Self::new();
        b.store.module_mut(b.module_addr).types = types;
        b
    }

    fn module_addr(&self) -> ModuleAddr {
        self.module_addr
    }

    fn store_mut(&mut self) -> &mut Store {
        &mut self.store
    }

    /// Register a wasm function body in the module.
    fn add_wasm(&mut self, ty: FuncType, locals: Vec<ValType>, body: Vec<Instruction>) -> FuncAddr {
        let addr = self.store.alloc_func(FunctionInstance::Wasm {
            ty,
            module: self.module_addr,
            locals,
            body,
        });
        self.store
            .module_mut(self.module_addr)
            .func_addrs
            .push(addr);
        addr
    }

    /// Register a host function in the module.
    fn add_host(&mut self, ty: FuncType, callback: HostFunc) -> FuncAddr {
        let addr = self
            .store
            .alloc_func(FunctionInstance::Host { ty, callback });
        self.store
            .module_mut(self.module_addr)
            .func_addrs
            .push(addr);
        addr
    }

    /// Register a funcref table backed by the module's function addresses.
    fn add_funcref_table(&mut self, entries: Vec<Ref>) -> we_wasm::runtime::TableAddr {
        let mut table = Table::new(TableType {
            elem: RefType::FuncRef,
            limits: Limits {
                min: entries.len() as u32,
                max: None,
            },
        })
        .unwrap();
        for (i, r) in entries.iter().enumerate() {
            table.set(i as u32, *r).unwrap();
        }
        let addr = self.store.alloc_table(table);
        self.store
            .module_mut(self.module_addr)
            .table_addrs
            .push(addr);
        addr
    }
}

fn ft(params: Vec<ValType>, results: Vec<ValType>) -> FuncType {
    FuncType { params, results }
}

// ---------------------------------------------------------------------------
// Basic invoke and arity
// ---------------------------------------------------------------------------

#[test]
fn invoke_constant_returning_function() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![Instruction::I32Const(42), Instruction::End],
    );
    let mut engine = Engine::new(b.store_mut());
    let out = engine.invoke(f, vec![]).unwrap();
    assert_eq!(out, vec![WasmValue::I32(42)]);
}

#[test]
fn invoke_passes_arguments_into_locals() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32, ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::LocalGet(1),
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    let out = engine
        .invoke(f, vec![WasmValue::I32(2), WasmValue::I32(40)])
        .unwrap();
    assert_eq!(out, vec![WasmValue::I32(42)]);
}

#[test]
fn invoke_initializes_declared_locals_to_zero() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![ValType::I32],
        vec![Instruction::LocalGet(0), Instruction::End],
    );
    let mut engine = Engine::new(b.store_mut());
    let out = engine.invoke(f, vec![]).unwrap();
    assert_eq!(out, vec![WasmValue::I32(0)]);
}

#[test]
fn arity_mismatch_traps() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![]),
        vec![],
        vec![Instruction::End],
    );
    let mut engine = Engine::new(b.store_mut());
    let err = engine.invoke(f, vec![]).unwrap_err();
    assert_eq!(err, Trap::IndirectCallTypeMismatch);
}

#[test]
fn type_mismatch_traps() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![]),
        vec![],
        vec![Instruction::End],
    );
    let mut engine = Engine::new(b.store_mut());
    let err = engine.invoke(f, vec![WasmValue::I64(0)]).unwrap_err();
    assert_eq!(err, Trap::IndirectCallTypeMismatch);
}

// ---------------------------------------------------------------------------
// Block / Loop / If / Else
// ---------------------------------------------------------------------------

#[test]
fn empty_block_is_no_op() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Empty),
            Instruction::End,
            Instruction::I32Const(7),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(7)]);
}

#[test]
fn block_returns_single_value() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::I32Const(100),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(100)]);
}

#[test]
fn br_skips_remaining_block_body() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::I32Const(1),
            Instruction::Br(0),
            Instruction::Drop,
            Instruction::I32Const(2),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(1)]);
}

#[test]
fn br_if_taken_and_not_taken() {
    let mut b = ModuleBuilder::new();
    // (param i32) -> i32. If param != 0 return 11 else return 22.
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::I32Const(11),
            Instruction::LocalGet(0),
            Instruction::BrIf(0),
            Instruction::Drop,
            Instruction::I32Const(22),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(1)]).unwrap(),
        vec![WasmValue::I32(11)]
    );
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(0)]).unwrap(),
        vec![WasmValue::I32(22)]
    );
}

#[test]
fn deeply_nested_blocks_break_through_levels() {
    let mut b = ModuleBuilder::new();
    // Three nested blocks, each with i32 result. Inner br 2 carries top
    // value all the way out.
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::I32Const(99),
            Instruction::Br(2),
            Instruction::End,
            Instruction::End,
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(99)]);
}

#[test]
fn if_then_branch_runs() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::If(BlockType::Value(ValType::I32)),
            Instruction::I32Const(10),
            Instruction::Else,
            Instruction::I32Const(20),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(1)]).unwrap(),
        vec![WasmValue::I32(10)]
    );
}

#[test]
fn if_else_branch_runs() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::If(BlockType::Value(ValType::I32)),
            Instruction::I32Const(10),
            Instruction::Else,
            Instruction::I32Const(20),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(0)]).unwrap(),
        vec![WasmValue::I32(20)]
    );
}

#[test]
fn if_without_else_with_false_condition_leaves_stack_unchanged() {
    let mut b = ModuleBuilder::new();
    // (param i32) -> i32. push 5, if local: (nothing), end, return 5.
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::I32Const(5),
            Instruction::LocalGet(0),
            Instruction::If(BlockType::Empty),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(0)]).unwrap(),
        vec![WasmValue::I32(5)]
    );
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(1)]).unwrap(),
        vec![WasmValue::I32(5)]
    );
}

#[test]
fn loop_with_br_iterates() {
    let mut b = ModuleBuilder::new();
    // Sum n down to 0: locals[0] = n, locals[1] = acc.
    // (param i32) -> i32 result is acc.
    // loop:
    //   if local[0] == 0 break out
    //   acc += local[0]
    //   local[0] -= 1
    //   br 0 (continue loop)
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![ValType::I32],
        vec![
            Instruction::Block(BlockType::Empty),
            Instruction::Loop(BlockType::Empty),
            // if local[0] == 0 -> break out
            Instruction::LocalGet(0),
            Instruction::I32Eqz,
            Instruction::BrIf(1),
            // acc += local[0]
            Instruction::LocalGet(1),
            Instruction::LocalGet(0),
            Instruction::I32Add,
            Instruction::LocalSet(1),
            // local[0] -= 1
            Instruction::LocalGet(0),
            Instruction::I32Const(1),
            Instruction::I32Sub,
            Instruction::LocalSet(0),
            // continue loop
            Instruction::Br(0),
            Instruction::End,
            Instruction::End,
            Instruction::LocalGet(1),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    // 5+4+3+2+1 = 15.
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(5)]).unwrap(),
        vec![WasmValue::I32(15)]
    );
}

// ---------------------------------------------------------------------------
// Multi-value blocks
// ---------------------------------------------------------------------------

#[test]
fn multi_value_block_returns_two_results() {
    let mut b = ModuleBuilder::with_types(vec![
        // type 0: () -> (i32 i32)
        ft(vec![], vec![ValType::I32, ValType::I32]),
    ]);
    // function: () -> i32, returns 100 + 200.
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::TypeIndex(0)),
            Instruction::I32Const(100),
            Instruction::I32Const(200),
            Instruction::End,
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(300)]);
}

#[test]
fn multi_value_block_with_params_and_results() {
    let mut b = ModuleBuilder::with_types(vec![
        // type 0: (i32 i32) -> (i32 i32) — swap-ish.
        ft(
            vec![ValType::I32, ValType::I32],
            vec![ValType::I32, ValType::I32],
        ),
    ]);
    // Push (3, 4), enter block taking 2 i32 returning 2 i32 (a, b -> b, a),
    // then sub.
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![ValType::I32],
        vec![
            Instruction::I32Const(3),
            Instruction::I32Const(4),
            Instruction::Block(BlockType::TypeIndex(0)),
            // a, b on stack. Want b, a. Save a, leave b, push a.
            Instruction::LocalSet(0),
            // stack: a, locals[0]=b. Wait, locals[0] is the swap slot.
            // Actually re-arrange: pop b into local[0], pop a -> push a, push b.
            // ugh, let me redo.
            // After Block: stack top is b, below is a.
            // local.set 0: pops b -> local[0] = b. stack: a.
            // We want stack ending as: b, a (b on top).
            // push local[0]? No, we want: stack: a, then b -> top b.
            // Actually current is: stack: a. We can: local.get 0 (push b)
            // -> stack: a, b. But we want b, a. Swap.
            // Tee b into local 0 first then swap?
            // Simpler: pop b into local 0, leave a, push local 0.
            // But the result is (a, b), not (b, a). Whatever — let me just test the multi-value mechanic.
            Instruction::LocalGet(0),
            Instruction::End,
            // Stack: a, b (i.e. block returned a, b on top). Sub: a-b = 3-4 = -1.
            Instruction::I32Sub,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(-1)]);
}

#[test]
fn multi_value_branch_carries_all_results() {
    let mut b = ModuleBuilder::with_types(vec![
        // type 0: () -> (i32 i32)
        ft(vec![], vec![ValType::I32, ValType::I32]),
    ]);
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::TypeIndex(0)),
            Instruction::I32Const(30),
            Instruction::I32Const(12),
            Instruction::Br(0),
            Instruction::I32Const(0),
            Instruction::I32Const(0),
            Instruction::End,
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(42)]);
}

// ---------------------------------------------------------------------------
// br_table
// ---------------------------------------------------------------------------

#[test]
fn br_table_picks_correct_target() {
    // Three blocks; param selects which one to break out of.
    // result is 100 / 200 / 300 depending on input 0 / 1 / >=2.
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::Block(BlockType::Empty),
            Instruction::Block(BlockType::Empty),
            Instruction::Block(BlockType::Empty),
            Instruction::LocalGet(0),
            Instruction::BrTable {
                labels: vec![0, 1, 2],
                default: 2,
            },
            Instruction::End, // closes innermost (label 0)
            Instruction::I32Const(100),
            Instruction::Br(2),
            Instruction::End, // closes (label 1)
            Instruction::I32Const(200),
            Instruction::Br(1),
            Instruction::End, // closes (label 2)
            Instruction::I32Const(300),
            Instruction::End, // closes outer i32 block
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(0)]).unwrap(),
        vec![WasmValue::I32(100)]
    );
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(1)]).unwrap(),
        vec![WasmValue::I32(200)]
    );
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(2)]).unwrap(),
        vec![WasmValue::I32(300)]
    );
    // Out-of-range -> default (label 2 -> 300).
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(99)]).unwrap(),
        vec![WasmValue::I32(300)]
    );
}

#[test]
fn br_table_with_carry_value() {
    // Two i32 result blocks; carry an i32 across the br_table.
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::Block(BlockType::Value(ValType::I32)),
            Instruction::I32Const(7),
            Instruction::LocalGet(0),
            Instruction::BrTable {
                labels: vec![0, 1],
                default: 1,
            },
            Instruction::End, // inner: pop 7 from carry, push 8
            Instruction::Drop,
            Instruction::I32Const(8),
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    // Input 0 -> branch to inner (label 0) carrying 7, push 8 -> return 8.
    // Wait — branching to label 0 (inner block result arity 1) — value on stack
    // after branch is 7. Then we hit End of inner block immediately and exit.
    // The inner produces result 7. Outer takes 7. Wait the layout is:
    // outer { inner { 7 brtable; if br to 0 -> outer sees inner=7 result -> then outer-end is hit }
    //         after end-of-inner: drop, push 8 } outer-end
    // Inner produces (7) -> outer continues after inner-end with 7 on stack,
    // then drops, pushes 8, end, end.
    // So input 0 actually drops 7, pushes 8 -> returns 8.
    // Input 1 -> br 1 (outer), carrying 7. Outer's End sees stack with 7
    // -> returns 7.
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(0)]).unwrap(),
        vec![WasmValue::I32(8)]
    );
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(1)]).unwrap(),
        vec![WasmValue::I32(7)]
    );
}

// ---------------------------------------------------------------------------
// Return
// ---------------------------------------------------------------------------

#[test]
fn return_short_circuits_function() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::I32Const(7),
            Instruction::Return,
            Instruction::Drop,
            Instruction::I32Const(99),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(7)]);
}

#[test]
fn return_inside_nested_blocks() {
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::Block(BlockType::Empty),
            Instruction::Block(BlockType::Empty),
            Instruction::I32Const(42),
            Instruction::Return,
            Instruction::End,
            Instruction::End,
            Instruction::I32Const(0),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(engine.invoke(f, vec![]).unwrap(), vec![WasmValue::I32(42)]);
}

// ---------------------------------------------------------------------------
// Call (direct)
// ---------------------------------------------------------------------------

#[test]
fn call_invokes_in_module_function() {
    let mut b = ModuleBuilder::new();
    // callee: (i32, i32) -> i32, returns a + b.
    let callee = b.add_wasm(
        ft(vec![ValType::I32, ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::LocalGet(1),
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    let callee_index = b.store.module(b.module_addr()).func_addrs.len() as u32 - 1;
    let _ = callee;
    // caller: () -> i32, returns callee(3, 4) + 1.
    let caller = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::I32Const(3),
            Instruction::I32Const(4),
            Instruction::Call(callee_index),
            Instruction::I32Const(1),
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(caller, vec![]).unwrap(),
        vec![WasmValue::I32(8)]
    );
}

#[test]
fn mutual_recursion_across_two_functions() {
    let mut b = ModuleBuilder::new();
    // Two functions:
    //   is_even(n): n == 0 ? true : is_odd(n-1)
    //   is_odd(n):  n == 0 ? false : is_even(n-1)
    // We pre-allocate addresses for both via indices 0 and 1.
    // To do that we add stubs first, then overwrite their bodies.
    let stub_ty = ft(vec![ValType::I32], vec![ValType::I32]);

    // Reserve func index 0 -> is_even, index 1 -> is_odd.
    let f0 = b.add_wasm(stub_ty.clone(), vec![], vec![Instruction::End]);
    let f1 = b.add_wasm(stub_ty.clone(), vec![], vec![Instruction::End]);

    // is_even body: if local.0 == 0 then 1 else call is_odd(local.0 - 1)
    let is_even_body = vec![
        Instruction::LocalGet(0),
        Instruction::I32Eqz,
        Instruction::If(BlockType::Value(ValType::I32)),
        Instruction::I32Const(1),
        Instruction::Else,
        Instruction::LocalGet(0),
        Instruction::I32Const(1),
        Instruction::I32Sub,
        Instruction::Call(1),
        Instruction::End,
        Instruction::End,
    ];
    let is_odd_body = vec![
        Instruction::LocalGet(0),
        Instruction::I32Eqz,
        Instruction::If(BlockType::Value(ValType::I32)),
        Instruction::I32Const(0),
        Instruction::Else,
        Instruction::LocalGet(0),
        Instruction::I32Const(1),
        Instruction::I32Sub,
        Instruction::Call(0),
        Instruction::End,
        Instruction::End,
    ];

    // Overwrite the function bodies.
    *b.store_mut().func_mut(f0) = FunctionInstance::Wasm {
        ty: stub_ty.clone(),
        module: b.module_addr(),
        locals: vec![],
        body: is_even_body,
    };
    *b.store_mut().func_mut(f1) = FunctionInstance::Wasm {
        ty: stub_ty,
        module: b.module_addr(),
        locals: vec![],
        body: is_odd_body,
    };

    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f0, vec![WasmValue::I32(10)]).unwrap(),
        vec![WasmValue::I32(1)],
        "10 is even"
    );
    assert_eq!(
        engine.invoke(f0, vec![WasmValue::I32(7)]).unwrap(),
        vec![WasmValue::I32(0)],
        "7 is not even"
    );
}

// ---------------------------------------------------------------------------
// Call indirect
// ---------------------------------------------------------------------------

#[test]
fn call_indirect_dispatches_to_table_entry() {
    let mut b = ModuleBuilder::with_types(vec![ft(vec![ValType::I32], vec![ValType::I32])]);
    // f0(n): n + 10
    let f0 = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::I32Const(10),
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    // f1(n): n + 20
    let f1 = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::I32Const(20),
            Instruction::I32Add,
            Instruction::End,
        ],
    );
    b.add_funcref_table(vec![Ref::Func(f0), Ref::Func(f1)]);

    // dispatch(i, n): call_indirect type 0 with operand n then i.
    let dispatch = b.add_wasm(
        ft(vec![ValType::I32, ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(1), // n
            Instruction::LocalGet(0), // i
            Instruction::CallIndirect {
                type_index: 0,
                table_index: 0,
            },
            Instruction::End,
        ],
    );

    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine
            .invoke(dispatch, vec![WasmValue::I32(0), WasmValue::I32(5)])
            .unwrap(),
        vec![WasmValue::I32(15)]
    );
    assert_eq!(
        engine
            .invoke(dispatch, vec![WasmValue::I32(1), WasmValue::I32(5)])
            .unwrap(),
        vec![WasmValue::I32(25)]
    );
}

#[test]
fn call_indirect_traps_on_null_funcref() {
    let mut b = ModuleBuilder::with_types(vec![ft(vec![], vec![])]);
    b.add_funcref_table(vec![Ref::Null(RefType::FuncRef)]);
    let f = b.add_wasm(
        ft(vec![], vec![]),
        vec![],
        vec![
            Instruction::I32Const(0),
            Instruction::CallIndirect {
                type_index: 0,
                table_index: 0,
            },
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![]).unwrap_err(),
        Trap::UninitializedElement
    );
}

#[test]
fn call_indirect_traps_on_type_mismatch() {
    // Table has a function with signature (i32) -> i32, but we call it with type 0 = () -> ().
    let mut b = ModuleBuilder::with_types(vec![ft(vec![], vec![])]);
    let target = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![Instruction::LocalGet(0), Instruction::End],
    );
    b.add_funcref_table(vec![Ref::Func(target)]);
    let f = b.add_wasm(
        ft(vec![], vec![]),
        vec![],
        vec![
            Instruction::I32Const(0),
            Instruction::CallIndirect {
                type_index: 0,
                table_index: 0,
            },
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![]).unwrap_err(),
        Trap::IndirectCallTypeMismatch
    );
}

#[test]
fn call_indirect_traps_on_out_of_bounds() {
    let mut b = ModuleBuilder::with_types(vec![ft(vec![], vec![])]);
    b.add_funcref_table(vec![]);
    let f = b.add_wasm(
        ft(vec![], vec![]),
        vec![],
        vec![
            Instruction::I32Const(0),
            Instruction::CallIndirect {
                type_index: 0,
                table_index: 0,
            },
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(f, vec![]).unwrap_err(),
        Trap::OutOfBoundsTable
    );
}

// ---------------------------------------------------------------------------
// Host function calls
// ---------------------------------------------------------------------------

#[test]
fn call_host_function_via_direct_call() {
    let mut b = ModuleBuilder::new();
    let host: HostFunc = Box::new(|args: &[WasmValue]| match args {
        [WasmValue::I32(a), WasmValue::I32(b)] => Ok(vec![WasmValue::I32(a * b)]),
        _ => unreachable!(),
    });
    let host_addr = b.add_host(
        ft(vec![ValType::I32, ValType::I32], vec![ValType::I32]),
        host,
    );
    let host_idx = b.store.module(b.module_addr()).func_addrs.len() as u32 - 1;
    let _ = host_addr;
    let caller = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::I32Const(6),
            Instruction::I32Const(7),
            Instruction::Call(host_idx),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(caller, vec![]).unwrap(),
        vec![WasmValue::I32(42)]
    );
}

#[test]
fn host_function_trap_propagates() {
    let mut b = ModuleBuilder::new();
    let host: HostFunc = Box::new(|_args: &[WasmValue]| Err(Trap::Unreachable));
    let host_addr = b.add_host(ft(vec![], vec![]), host);
    let host_idx = b.store.module(b.module_addr()).func_addrs.len() as u32 - 1;
    let _ = host_addr;
    let caller = b.add_wasm(
        ft(vec![], vec![]),
        vec![],
        vec![Instruction::Call(host_idx), Instruction::End],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(caller, vec![]).unwrap_err(),
        Trap::Unreachable
    );
}

#[test]
fn call_indirect_lands_on_host_function() {
    let mut b = ModuleBuilder::with_types(vec![ft(vec![ValType::I32], vec![ValType::I32])]);
    let host: HostFunc = Box::new(|args: &[WasmValue]| match args {
        [WasmValue::I32(a)] => Ok(vec![WasmValue::I32(a + 100)]),
        _ => unreachable!(),
    });
    let host_addr = b.add_host(ft(vec![ValType::I32], vec![ValType::I32]), host);
    b.add_funcref_table(vec![Ref::Func(host_addr)]);
    let caller = b.add_wasm(
        ft(vec![], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::I32Const(5),
            Instruction::I32Const(0),
            Instruction::CallIndirect {
                type_index: 0,
                table_index: 0,
            },
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    assert_eq!(
        engine.invoke(caller, vec![]).unwrap(),
        vec![WasmValue::I32(105)]
    );
}

// ---------------------------------------------------------------------------
// Stack overflow trap
// ---------------------------------------------------------------------------

#[test]
fn deep_recursion_traps_with_configurable_limit() {
    let mut b = ModuleBuilder::new();
    // A function that just calls itself.
    let f = b.add_wasm(
        ft(vec![], vec![]),
        vec![],
        vec![Instruction::Call(0), Instruction::End],
    );
    let mut engine = Engine::with_max_depth(b.store_mut(), 16);
    assert_eq!(engine.invoke(f, vec![]).unwrap_err(), Trap::StackOverflow);
}

#[test]
fn recursion_within_limit_succeeds() {
    let mut b = ModuleBuilder::new();
    // Sum from n down to 0 recursively. Avoid running too deep so this stays
    // within the default limit.
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![],
        vec![
            Instruction::LocalGet(0),
            Instruction::I32Eqz,
            Instruction::If(BlockType::Value(ValType::I32)),
            Instruction::I32Const(0),
            Instruction::Else,
            Instruction::LocalGet(0),
            Instruction::LocalGet(0),
            Instruction::I32Const(1),
            Instruction::I32Sub,
            Instruction::Call(0),
            Instruction::I32Add,
            Instruction::End,
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    // sum 1..=20 = 210.
    assert_eq!(
        engine.invoke(f, vec![WasmValue::I32(20)]).unwrap(),
        vec![WasmValue::I32(210)]
    );
}

// ---------------------------------------------------------------------------
// Integration: function body uses straight-line ops AND control flow together.
// ---------------------------------------------------------------------------

#[test]
fn fibonacci_via_loop() {
    // Iterative fib(n): a=0, b=1; for i in 0..n: t=a+b; a=b; b=t; -> a.
    let mut b = ModuleBuilder::new();
    let f = b.add_wasm(
        ft(vec![ValType::I32], vec![ValType::I32]),
        vec![ValType::I32, ValType::I32, ValType::I32, ValType::I32],
        // locals: [n (param), a, b, i, tmp]
        vec![
            // a = 0
            Instruction::I32Const(0),
            Instruction::LocalSet(1),
            // b = 1
            Instruction::I32Const(1),
            Instruction::LocalSet(2),
            // i = 0
            Instruction::I32Const(0),
            Instruction::LocalSet(3),
            Instruction::Block(BlockType::Empty),
            Instruction::Loop(BlockType::Empty),
            // if i >= n -> exit
            Instruction::LocalGet(3),
            Instruction::LocalGet(0),
            Instruction::I32GeS,
            Instruction::BrIf(1),
            // tmp = a + b
            Instruction::LocalGet(1),
            Instruction::LocalGet(2),
            Instruction::I32Add,
            Instruction::LocalSet(4),
            // a = b
            Instruction::LocalGet(2),
            Instruction::LocalSet(1),
            // b = tmp
            Instruction::LocalGet(4),
            Instruction::LocalSet(2),
            // i += 1
            Instruction::LocalGet(3),
            Instruction::I32Const(1),
            Instruction::I32Add,
            Instruction::LocalSet(3),
            Instruction::Br(0),
            Instruction::End,
            Instruction::End,
            Instruction::LocalGet(1),
            Instruction::End,
        ],
    );
    let mut engine = Engine::new(b.store_mut());
    // fib(0)=0, fib(1)=1, fib(10)=55, fib(20)=6765
    for (n, expected) in [(0, 0), (1, 1), (2, 1), (10, 55), (20, 6765)] {
        assert_eq!(
            engine.invoke(f, vec![WasmValue::I32(n)]).unwrap(),
            vec![WasmValue::I32(expected)],
            "fib({n})"
        );
    }
}