#include <z3++.h>

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

#include "fold/ir/call.hpp"
#include "fold/ir/expr.hpp"
#include "fold/ir/source_loc.hpp"
#include "fold/ir/symbol.hpp"
#include "fold/symbolic/safety_observer.hpp"

using namespace fold;
using namespace fold::ir;

// Minimal observer that records the ir-based hook calls.
class RecordingObserver : public SafetyObserver {
  public:
    // ir-based hook overrides
    void onFunctionStartIr(const std::string& function_name, z3::solver* solver) override {
        last_fn_start = function_name;
        fn_start_count++;
        (void)solver;
    }

    void onFunctionEndIr(const std::string& function_name) override {
        last_fn_end = function_name;
        fn_end_count++;
    }

    void onPointerDerefIr(const z3::expr& ptr, const ir::Expr& ptrExpr, const z3::expr& path,
                          ir::Symbol var_sym, const ir::SourceLoc& loc,
                          const symbolic::HeapAccessInfo* heap_access, const DefUseContext* def_use,
                          const SymbolicEnv* env) override {
        last_deref_sym = var_sym;
        deref_count++;
        (void)ptr;
        (void)ptrExpr;
        (void)path;
        (void)loc;
        (void)heap_access;
        (void)def_use;
        (void)env;
    }

    void onAssignmentIr(ir::Symbol lhs_sym, const ir::Expr& rhs_expr, const z3::expr& rhs_value,
                        const z3::expr& assigned_value, const ir::SourceLoc& loc) override {
        last_assign_sym = lhs_sym;
        assign_count++;
        (void)rhs_expr;
        (void)rhs_value;
        (void)assigned_value;
        (void)loc;
    }

    void onCallExprIr(const std::string& callee_name, const std::vector<ir::Expr>& arg_exprs,
                      const std::vector<z3::expr>& arg_values, const z3::expr& path,
                      const ir::SourceLoc& loc, const SymbolicEnv* env) override {
        last_callee = callee_name;
        call_count++;
        (void)arg_exprs;
        (void)arg_values;
        (void)path;
        (void)loc;
        (void)env;
    }

    void onArraySubscriptIr(const z3::expr& idx, int known_len, const z3::expr& path,
                            const ir::SourceLoc& loc, ir::Expr idx_expr,
                            const symbolic::HeapAccessInfo* heap_access,
                            const DefUseContext* def_use, const SymbolicEnv* env) override {
        last_known_len = known_len;
        subscript_count++;
        (void)idx;
        (void)path;
        (void)loc;
        (void)idx_expr;
        (void)heap_access;
        (void)def_use;
        (void)env;
    }

    void onMemberAccessIr(ir::Expr base_expr, const std::string& field_name,
                          const z3::expr& base_value, const z3::expr& result, const z3::expr& path,
                          const ir::SourceLoc& loc) override {
        last_field = field_name;
        member_count++;
        (void)base_expr;
        (void)base_value;
        (void)result;
        (void)path;
        (void)loc;
    }

    // Counters and last-seen values.
    int fn_start_count = 0;
    int fn_end_count = 0;
    int deref_count = 0;
    int assign_count = 0;
    int call_count = 0;
    int subscript_count = 0;
    int member_count = 0;

    std::string last_fn_start;
    std::string last_fn_end;
    ir::Symbol last_deref_sym;
    ir::Symbol last_assign_sym;
    std::string last_callee;
    std::string last_field;
    int last_known_len = -1;
};

TEST_CASE("SafetyObserver: ir-based onFunctionStartIr called with correct name", "[observer_ir]") {
    RecordingObserver obs;
    z3::context ctx;
    z3::solver solver(ctx, "QF_BV");

    obs.onFunctionStartIr("crate::module::func", &solver);

    REQUIRE(obs.fn_start_count == 1);
    REQUIRE(obs.last_fn_start == "crate::module::func");
}

TEST_CASE("SafetyObserver: ir-based onFunctionEndIr called with correct name", "[observer_ir]") {
    RecordingObserver obs;
    obs.onFunctionEndIr("crate::module::func");

    REQUIRE(obs.fn_end_count == 1);
    REQUIRE(obs.last_fn_end == "crate::module::func");
}

TEST_CASE("SafetyObserver: ir-based onPointerDerefIr works without clang stubs", "[observer_ir]") {
    RecordingObserver obs;

    z3::context ctx;
    z3::expr ptr = ctx.bv_val(0x1000, 64);
    ir::Expr ptrExpr;
    z3::expr path = ctx.bool_val(true);
    ir::Symbol sym(reinterpret_cast<const void*>(0xDEADBEEF));
    ir::SourceLoc loc;

    obs.onPointerDerefIr(ptr, ptrExpr, path, sym, loc, nullptr, nullptr, nullptr);

    REQUIRE(obs.deref_count == 1);
    REQUIRE(obs.last_deref_sym == sym);
}

TEST_CASE("SafetyObserver: ir-based onAssignmentIr works without clang stubs", "[observer_ir]") {
    RecordingObserver obs;

    z3::context ctx;
    ir::Symbol sym(reinterpret_cast<const void*>(0xCAFE));
    ir::Expr rhs;
    z3::expr rhs_val = ctx.bv_val(42, 64);
    z3::expr assigned = ctx.bv_val(42, 64);
    ir::SourceLoc loc;

    obs.onAssignmentIr(sym, rhs, rhs_val, assigned, loc);

    REQUIRE(obs.assign_count == 1);
    REQUIRE(obs.last_assign_sym == sym);
}

TEST_CASE("SafetyObserver: ir-based onCallExprIr works without clang stubs", "[observer_ir]") {
    RecordingObserver obs;

    z3::context ctx;
    std::string callee = "std::vec::Vec::push";
    std::vector<ir::Expr> arg_exprs;
    std::vector<z3::expr> arg_vals;
    z3::expr path = ctx.bool_val(true);
    ir::SourceLoc loc;

    obs.onCallExprIr(callee, arg_exprs, arg_vals, path, loc, nullptr);

    REQUIRE(obs.call_count == 1);
    REQUIRE(obs.last_callee == "std::vec::Vec::push");
}

TEST_CASE("SafetyObserver: ir-based onArraySubscriptIr records known_len", "[observer_ir]") {
    RecordingObserver obs;

    z3::context ctx;
    z3::expr idx = ctx.bv_val(5, 64);
    z3::expr path = ctx.bool_val(true);
    ir::SourceLoc loc;
    ir::Expr idx_expr;

    obs.onArraySubscriptIr(idx, 100, path, loc, idx_expr, nullptr, nullptr, nullptr);

    REQUIRE(obs.subscript_count == 1);
    REQUIRE(obs.last_known_len == 100);
}

TEST_CASE("SafetyObserver: ir-based onMemberAccessIr records field name", "[observer_ir]") {
    RecordingObserver obs;

    z3::context ctx;
    ir::Expr base;
    std::string field = "data";
    z3::expr base_val = ctx.bv_val(0, 64);
    z3::expr result = ctx.bv_val(42, 64);
    z3::expr path = ctx.bool_val(true);
    ir::SourceLoc loc;

    obs.onMemberAccessIr(base, field, base_val, result, path, loc);

    REQUIRE(obs.member_count == 1);
    REQUIRE(obs.last_field == "data");
}

TEST_CASE("SafetyObserver: default implementations are no-ops", "[observer_ir]") {
    // The base SafetyObserver itself should not crash on any ir-based hook.
    SafetyObserver obs;

    z3::context ctx;
    z3::solver solver(ctx, "QF_BV");
    z3::expr dummy = ctx.bv_val(0, 64);
    ir::Expr dummy_expr;
    z3::expr path = ctx.bool_val(true);
    ir::Symbol sym;
    ir::SourceLoc loc;
    std::vector<ir::Expr> arg_exprs;
    std::vector<z3::expr> arg_vals;

    // All default calls should be no-ops (no crash).
    obs.onFunctionStartIr("test", &solver);
    obs.onFunctionEndIr("test");
    obs.onPointerDerefIr(dummy, dummy_expr, path, sym, loc, nullptr, nullptr, nullptr);
    obs.onAssignmentIr(sym, dummy_expr, dummy, dummy, loc);
    obs.onCallExprIr("test", arg_exprs, arg_vals, path, loc, nullptr);
    obs.onArraySubscriptIr(dummy, -1, path, loc, dummy_expr, nullptr, nullptr, nullptr);
    obs.onMemberAccessIr(dummy_expr, "f", dummy, dummy, path, loc);

    SUCCEED("Default implementations did not crash");
}