#include <clang/AST/RecursiveASTVisitor.h>
#include <clang/AST/Stmt.h>
#include <clang/Frontend/ASTUnit.h>
#include <clang/Tooling/Tooling.h>

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

#include "fold/symbolic/assumption_provenance.hpp"
#include "fold/symbolic/encoder.hpp"
#include "fold/symbolic/loop_bound_pin.hpp"
#include "fold/symbolic/modeling_options.hpp"

using fold::ModelingOptions;
using fold::SymbolicEnv;

namespace {

// Grab the function body's DeclStmt (`int i = 0;`) and the first IfStmt's
// condition (`i < 6`) from a parsed snippet.
struct Grabber : clang::RecursiveASTVisitor<Grabber> {
    const clang::DeclStmt* decl = nullptr;
    const clang::Expr* cond = nullptr;
    bool VisitDeclStmt(clang::DeclStmt* ds) {
        if (!decl)
            decl = ds;
        return true;
    }
    bool VisitIfStmt(clang::IfStmt* is) {
        if (!cond)
            cond = is->getCond();
        return true;
    }
};

std::unique_ptr<clang::ASTUnit> parse(const std::string& code) {
    return clang::tooling::buildASTFromCodeWithArgs(code, {"-std=c++20"}, "input.cpp");
}

// Returns {tracked-assumption count, binding-is-numeral} after applying the
// DeclStmt (i=0) and running pinLoopBoundVars over the `i < 6` condition.
struct PinOutcome {
    std::size_t tracked = 0;
    bool numeral = true;
    bool first_is_loop_bound = false;
};

PinOutcome runPin(clang::ASTUnit& unit, const ModelingOptions& m) {
    auto& ctx = unit.getASTContext();
    Grabber g;
    g.TraverseDecl(ctx.getTranslationUnitDecl());
    REQUIRE(g.decl);
    REQUIRE(g.cond);
    const auto* vd = clang::dyn_cast<clang::VarDecl>(*g.decl->decl_begin())->getCanonicalDecl();

    SymbolicEnv env = SymbolicEnv::create();
    env.setModelingOptions(m);
    fold::applyStmt(g.decl, ctx, env); // binds i to the numeral 0
    REQUIRE(env.tryGet(vd).has_value());
    REQUIRE(env.tryGet(vd)->is_numeral());

    fold::symbolic::pinLoopBoundVars(g.cond, env);

    PinOutcome out;
    out.tracked = env.trackedAssumptions().size();
    out.numeral = env.tryGet(vd)->is_numeral();
    if (!env.trackedAssumptions().empty())
        out.first_is_loop_bound = env.trackedAssumptions().front().prov.kind ==
                                  fold::symbolic::AssumptionKind::LoopBoundUnroll;
    return out;
}

constexpr const char* kCode = "void f(){ int i = 0; if (i < 6) { } }";

} // namespace

TEST_CASE("pinLoopBoundVars pins a concrete loop var under the toggle", "[loop-bound-pin]") {
    auto unit = parse(kCode);
    REQUIRE(unit);
    ModelingOptions on;
    on.suppress_loop_bound_always_conditions = true;
    auto out = runPin(*unit, on);
    // The IV was rebound to a fresh symbolic, recorded as one tracked assumption.
    REQUIRE(out.tracked == 1);
    REQUIRE_FALSE(out.numeral);
    REQUIRE(out.first_is_loop_bound);
}

TEST_CASE("pinLoopBoundVars is a no-op when the toggle is off", "[loop-bound-pin]") {
    auto unit = parse(kCode);
    REQUIRE(unit);
    ModelingOptions off; // all false
    auto out = runPin(*unit, off);
    REQUIRE(out.tracked == 0);
    REQUIRE(out.numeral); // binding unchanged
}