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

#include "fold/symbolic/encoder.hpp"
#include "fold/symbolic/encoder_reachability.hpp"

namespace {

// Tests in this file mutate the process-wide reachability flag/counter.
// Wrap each case with a guard that restores the prior state.
class ProfilingGuard {
  public:
    ProfilingGuard()
        : was_enabled_(fold::reachabilityProfilingEnabled()),
          prior_encoded_(fold::globalReachabilityCounter().encoded),
          prior_bailed_(fold::globalReachabilityCounter().bailed),
          prior_reasons_(fold::globalReachabilityCounter().bail_reasons) {}
    ~ProfilingGuard() {
        fold::setReachabilityProfilingEnabled(was_enabled_);
        fold::globalReachabilityCounter().encoded = std::move(prior_encoded_);
        fold::globalReachabilityCounter().bailed = std::move(prior_bailed_);
        fold::globalReachabilityCounter().bail_reasons = std::move(prior_reasons_);
    }
    ProfilingGuard(const ProfilingGuard&) = delete;
    ProfilingGuard& operator=(const ProfilingGuard&) = delete;

  private:
    bool was_enabled_;
    decltype(fold::ReachabilityCounter::encoded) prior_encoded_;
    decltype(fold::ReachabilityCounter::bailed) prior_bailed_;
    decltype(fold::ReachabilityCounter::bail_reasons) prior_reasons_;
};

} // namespace

TEST_CASE("markEncoded is a no-op on counters when profiling is disabled",
          "[encoder][reachability]") {
    ProfilingGuard guard;
    fold::setReachabilityProfilingEnabled(false);
    fold::resetReachabilityCounter();

    z3::context ctx;
    z3::expr value = ctx.bv_val(42, fold::kBitvecWidth);
    z3::expr wrapped = fold::markEncoded("TestNode", value);

    REQUIRE(wrapped.is_bv());
    REQUIRE(fold::globalReachabilityCounter().encoded.empty());
}

TEST_CASE("markEncoded records the node kind when profiling is enabled",
          "[encoder][reachability]") {
    ProfilingGuard guard;
    fold::setReachabilityProfilingEnabled(true);
    fold::resetReachabilityCounter();

    z3::context ctx;
    z3::expr value = ctx.bool_val(true);
    z3::expr a = fold::markEncoded("BinaryOperator", value);
    z3::expr b = fold::markEncoded("BinaryOperator", value);
    z3::expr c = fold::markEncoded("CallExpr", value);

    const auto& counter = fold::globalReachabilityCounter();
    REQUIRE(counter.encoded.at("BinaryOperator") == 2);
    REQUIRE(counter.encoded.at("CallExpr") == 1);
    REQUIRE(counter.bailed.empty());
}

TEST_CASE("recordBailedNode records reason buckets", "[encoder][reachability]") {
    ProfilingGuard guard;
    fold::setReachabilityProfilingEnabled(true);
    fold::resetReachabilityCounter();

    fold::recordBailedNode("UnaryOperator", "unsupported opcode");
    fold::recordBailedNode("UnaryOperator", "unsupported opcode");
    fold::recordBailedNode("UnaryOperator", "missing operand");

    const auto& counter = fold::globalReachabilityCounter();
    REQUIRE(counter.bailed.at("UnaryOperator") == 3);
    REQUIRE(counter.bail_reasons.at("UnaryOperator").at("unsupported opcode") == 2);
    REQUIRE(counter.bail_reasons.at("UnaryOperator").at("missing operand") == 1);
}

TEST_CASE("resetReachabilityCounter clears all buckets", "[encoder][reachability]") {
    ProfilingGuard guard;
    fold::setReachabilityProfilingEnabled(true);
    fold::resetReachabilityCounter();
    fold::recordEncoded("DeclRef");
    fold::recordBailedNode("CallExpr", "external");

    fold::resetReachabilityCounter();

    const auto& counter = fold::globalReachabilityCounter();
    REQUIRE(counter.encoded.empty());
    REQUIRE(counter.bailed.empty());
    REQUIRE(counter.bail_reasons.empty());
}

TEST_CASE("ReachabilityCounter::print includes node kinds and bail rates",
          "[encoder][reachability]") {
    fold::ReachabilityCounter counter;
    counter.record_encoded("DeclRef");
    counter.record_encoded("DeclRef");
    counter.record_bailed("CallExpr", "external");

    std::ostringstream out;
    counter.print(out);
    const std::string s = out.str();

    REQUIRE(s.find("reachability") != std::string::npos);
    REQUIRE(s.find("DeclRef") != std::string::npos);
    REQUIRE(s.find("CallExpr") != std::string::npos);
    REQUIRE(s.find("external(1)") != std::string::npos);
}

TEST_CASE("setReachabilityProfilingEnabled flips the toggle", "[encoder][reachability]") {
    ProfilingGuard guard;
    fold::setReachabilityProfilingEnabled(false);
    REQUIRE_FALSE(fold::reachabilityProfilingEnabled());
    fold::setReachabilityProfilingEnabled(true);
    REQUIRE(fold::reachabilityProfilingEnabled());
}