#include <clang/AST/RecursiveASTVisitor.h>
#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/null_deref.hpp"
#include "fold/detect/summary_store.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"
#include "fold/symbolic/encoder_callee_names.hpp"
#include "fold/symbolic/summary.hpp"
#include "fold/symbolic/summary_encoder.hpp"

// Function summaries were stored under the qualified name, which overloads
// share: `take(int*, int)` (dereferences p) and `take(int*, long)` (checks p)
// both mapped to "take", so whichever was summarized first answered for both.
// A caller of the null-checking overload then inherited a nonnull
// precondition it does not have, and vice versa. Summaries are now keyed by
// the declaration's USR.

namespace {

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

std::vector<const clang::FunctionDecl*> definitionsNamed(clang::ASTUnit& unit,
                                                         const std::string& name) {
    struct Finder : clang::RecursiveASTVisitor<Finder> {
        std::string name;
        std::vector<const clang::FunctionDecl*> found;
        bool VisitFunctionDecl(clang::FunctionDecl* fd) {
            if (fd->getNameAsString() == name && fd->isThisDeclarationADefinition())
                found.push_back(fd);
            return true;
        }
    } finder;
    finder.name = name;
    finder.TraverseDecl(unit.getASTContext().getTranslationUnitDecl());
    return finder.found;
}

bool requiresNonNull(const fold::symbolic::SymbolicSummary* summary, unsigned index) {
    if (summary == nullptr)
        return false;
    const auto it = summary->param_pre.find(index);
    if (it == summary->param_pre.end())
        return false;
    for (const auto& term : it->second)
        if (term.kind == "nonnull")
            return true;
    return false;
}

} // namespace

TEST_CASE("overloads get separate summaries", "[summary-key]") {
    fold::detect::SummaryStore::instance().clear();
    auto unit = parse("int take(int* p, int) { return *p; }\n"
                      "int take(int* p, long) { return p ? *p : 0; }\n");
    REQUIRE(unit);
    const auto defs = definitionsNamed(*unit, "take");
    REQUIRE(defs.size() == 2);
    const auto& ctx = unit->getASTContext();

    SECTION("dereferencing overload first") {
        CHECK(requiresNonNull(fold::ensureSymbolicSummary(defs[0], ctx), 0));
        CHECK_FALSE(requiresNonNull(fold::ensureSymbolicSummary(defs[1], ctx), 0));
    }
    SECTION("null-checking overload first") {
        CHECK_FALSE(requiresNonNull(fold::ensureSymbolicSummary(defs[1], ctx), 0));
        CHECK(requiresNonNull(fold::ensureSymbolicSummary(defs[0], ctx), 0));
    }
    fold::detect::SummaryStore::instance().clear();
}

TEST_CASE("summary keys: identity across redeclarations, not across names", "[summary-key]") {
    auto unit = parse("int f(int* p);\n"
                      "int f(int* p) { return *p; }\n"
                      "namespace a { int g(int* p) { return *p; } }\n"
                      "namespace b { int g(int* p) { return p ? *p : 0; } }\n");
    REQUIRE(unit);
    struct Finder : clang::RecursiveASTVisitor<Finder> {
        std::vector<const clang::FunctionDecl*> f_decls;
        std::vector<const clang::FunctionDecl*> g_defs;
        bool VisitFunctionDecl(clang::FunctionDecl* fd) {
            if (fd->getNameAsString() == "f")
                f_decls.push_back(fd);
            if (fd->getNameAsString() == "g")
                g_defs.push_back(fd);
            return true;
        }
    } finder;
    finder.TraverseDecl(unit->getASTContext().getTranslationUnitDecl());
    REQUIRE(finder.f_decls.size() == 2);
    REQUIRE(finder.g_defs.size() == 2);
    CHECK(fold::summaryStoreKey(finder.f_decls[0]) == fold::summaryStoreKey(finder.f_decls[1]));
    CHECK(fold::summaryStoreKey(finder.g_defs[0]) != fold::summaryStoreKey(finder.g_defs[1]));
    CHECK_FALSE(fold::summaryStoreKey(finder.f_decls[0]).empty());
    CHECK(fold::summaryStoreKey(nullptr).empty());
}

TEST_CASE("null-deref: a cross-TU call uses its own overload's summary",
          "[summary-key][null-deref]") {
    // TU 1 defines both overloads and summarizes them (the dereferencing one
    // last). TU 2 sees only declarations, so its call resolves through the
    // store alone, and passes null to the overload that checks.
    const std::string defs = "int take(int* p, long) { return p ? *p : 0; }\n"
                             "int take(int* p, int) { return *p; }\n";
    const std::string caller = "int take(int* p, long);\n"
                               "int take(int* p, int);\n"
                               "int caller() { int* q = nullptr; return take(q, 1L); }\n"
                               "int control() { int* q = nullptr; return take(q, 1); }\n";
    fold::detect::SummaryStore::instance().clear();
    auto defs_unit = parse(defs);
    REQUIRE(defs_unit);
    for (const auto* fd : definitionsNamed(*defs_unit, "take"))
        (void)fold::ensureSymbolicSummary(fd, defs_unit->getASTContext());

    auto caller_unit = parse(caller);
    REQUIRE(caller_unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    const auto findings = fold::detectNullDerefInContext(caller_unit->getASTContext(), b, s);
    fold::detect::SummaryStore::instance().clear();
    std::size_t in_caller = 0;
    std::size_t in_control = 0;
    for (const auto& f : findings) {
        if (!f.locations.empty() && f.locations.front().line == 3)
            ++in_caller;
        if (!f.locations.empty() && f.locations.front().line == 4)
            ++in_control;
    }
    CHECK(in_caller == 0);
    // Control: the dereferencing overload's precondition still applies.
    CHECK(in_control == 1);
}