#include <clang/Frontend/ASTUnit.h>
#include <clang/Tooling/Tooling.h>

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

#include "fold/detect/equiv.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"
#include "fold/symbolic/serialize.hpp"
#include "fold/symbolic/summary_encoder.hpp"

namespace {

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

// Pick the FunctionDecl with the given name from a TU, serialize it.
fold::SerializedSummary serializeNamed(clang::ASTUnit& unit, const std::string& name,
                                       const std::string& file, int line) {
    auto& ctx = unit.getASTContext();
    for (auto* decl : ctx.getTranslationUnitDecl()->decls()) {
        if (auto* fd = clang::dyn_cast<clang::FunctionDecl>(decl)) {
            if (fd->getQualifiedNameAsString() == name && fd->hasBody()) {
                auto summary = fold::summarizeFunction(fd, ctx);
                REQUIRE(summary);
                std::string sig = fd->getReturnType().getCanonicalType().getAsString() + "(";
                bool first = true;
                for (const auto* p : fd->parameters()) {
                    if (!first)
                        sig += ",";
                    sig += p->getType().getCanonicalType().getAsString();
                    first = false;
                }
                sig += ")";
                return fold::serialize(name + "@" + file + ":" + std::to_string(line), name, file,
                                       line, 1, sig, *summary);
            }
        }
    }
    FAIL("function not found: " << name);
    return {};
}

} // namespace

TEST_CASE("cross-TU equiv: abs_a and abs_b in different ASTUnits", "[crosstu][equiv]") {
    auto unit_a = parse("int abs_a(int x){ if(x<0) return -x; return x; }");
    auto unit_b = parse("int abs_b(int y){ return y<0 ? -y : y; }");
    REQUIRE(unit_a);
    REQUIRE(unit_b);

    std::vector<fold::SerializedSummary> serialized;
    serialized.push_back(serializeNamed(*unit_a, "abs_a", "a.cpp", 1));
    serialized.push_back(serializeNamed(*unit_b, "abs_b", "b.cpp", 1));
    const auto restored = fold::fromJson(fold::toJson(serialized.front()));
    REQUIRE(restored.effects.isObservationallyPure());

    fold::SolverBudget b;
    fold::SolverStats s;
    s.profiling_enabled = true;
    auto findings = fold::detectEquivalentFunctionsCrossTU(serialized, b, s);
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::Equiv);
    REQUIRE(findings[0].confidence == fold::Confidence::Proved);
    REQUIRE(findings[0].locations.size() == 2);
    REQUIRE(findings[0].evidence.quality.has_value());
    REQUIRE(findings[0].evidence.quality->rank_reasons ==
            std::vector<std::string>{"return-proof", "effects-empty", "cross-tu"});
    REQUIRE(s.equiv_cross_tu_pairs_considered >= 1);
}

TEST_CASE("cross-TU equiv rejects equal returns with external writes",
          "[crosstu][equiv][effects][negative]") {
    auto unit_a = parse("int state_a; int set_a(int x){ state_a = 1; return x + 1; }");
    auto unit_b = parse("int state_b; int set_b(int x){ state_b = 2; return x + 1; }");
    REQUIRE(unit_a);
    REQUIRE(unit_b);
    std::vector<fold::SerializedSummary> serialized;
    serialized.push_back(serializeNamed(*unit_a, "set_a", "a.cpp", 1));
    serialized.push_back(serializeNamed(*unit_b, "set_b", "b.cpp", 1));

    fold::SolverBudget budget;
    fold::SolverStats stats;
    REQUIRE(fold::detectEquivalentFunctionsCrossTU(serialized, budget, stats).empty());
}

TEST_CASE("cross-TU equivalence fails closed when effect metadata is absent",
          "[crosstu][equiv][effects][wire]") {
    auto unit_a = parse("int first(int x){ return x + 1; }");
    auto unit_b = parse("int second(int x){ return x + 1; }");
    REQUIRE(unit_a);
    REQUIRE(unit_b);
    auto first = fold::toJson(serializeNamed(*unit_a, "first", "a.cpp", 1));
    auto second = fold::toJson(serializeNamed(*unit_b, "second", "b.cpp", 1));
    first.erase("effects");
    second.erase("effects");
    std::vector<fold::SerializedSummary> serialized{fold::fromJson(first), fold::fromJson(second)};

    fold::SolverBudget budget;
    fold::SolverStats stats;
    REQUIRE(fold::detectEquivalentFunctionsCrossTU(serialized, budget, stats).empty());
}

TEST_CASE("cross-TU equiv: identity vs +1 must NOT match", "[crosstu][equiv][negative]") {
    auto unit_a = parse("int id(int x){ return x; }");
    auto unit_b = parse("int off(int x){ return x + 1; }");
    REQUIRE(unit_a);
    REQUIRE(unit_b);

    std::vector<fold::SerializedSummary> serialized;
    serialized.push_back(serializeNamed(*unit_a, "id", "a.cpp", 1));
    serialized.push_back(serializeNamed(*unit_b, "off", "b.cpp", 1));

    fold::SolverBudget b;
    fold::SolverStats s;
    auto findings = fold::detectEquivalentFunctionsCrossTU(serialized, b, s);
    REQUIRE(findings.empty());
}

TEST_CASE("cross-TU equiv: signature mismatch never pairs", "[crosstu][equiv][sig]") {
    auto unit_a = parse("int unary(int x){ return x; }");
    auto unit_b = parse("int binary(int x, int y){ return x; }");
    REQUIRE(unit_a);
    REQUIRE(unit_b);

    std::vector<fold::SerializedSummary> serialized;
    serialized.push_back(serializeNamed(*unit_a, "unary", "a.cpp", 1));
    serialized.push_back(serializeNamed(*unit_b, "binary", "b.cpp", 1));

    fold::SolverBudget b;
    fold::SolverStats s;
    auto findings = fold::detectEquivalentFunctionsCrossTU(serialized, b, s);
    REQUIRE(findings.empty());
}