// Phase 2 (interprocedural taint): taint threads across user-function call
// boundaries — through return values (conduits) and into forwarding sinks —
// while domain-matched sanitization inside a helper still suppresses.
#include <clang/AST/Decl.h>
#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/taint_domain.hpp"
#include "fold/detect/taint_summary.hpp"
#include "fold/detect/taint_track.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"

namespace {

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

const clang::FunctionDecl* findFn(clang::ASTUnit& unit, const std::string& name) {
    struct V : clang::RecursiveASTVisitor<V> {
        std::string target;
        const clang::FunctionDecl* found = nullptr;
        bool VisitFunctionDecl(clang::FunctionDecl* fd) {
            if (fd->getNameAsString() == target && fd->doesThisDeclarationHaveABody())
                found = fd->getCanonicalDecl();
            return true;
        }
    } v;
    v.target = name;
    v.TraverseDecl(unit.getASTContext().getTranslationUnitDecl());
    return v.found;
}

std::size_t runTaint(clang::ASTUnit& unit) {
    fold::SolverBudget b;
    fold::SolverStats s;
    auto fs = fold::detectTaintTrackInContext(unit.getASTContext(), b, s);
    std::size_t n = 0;
    for (const auto& f : fs)
        if (f.kind == fold::FindingKind::TaintTrack)
            ++n;
    return n;
}

constexpr const char* kDecls =
    "typedef struct MYSQL MYSQL;\n"
    "extern \"C\" char* getenv(const char*);\n"
    "extern \"C\" int system(const char*);\n"
    "extern \"C\" int mysql_query(MYSQL*, const char*);\n"
    "extern \"C\" char* PQescapeLiteral(void*, const char*, unsigned long);\n";

} // namespace

TEST_CASE("summary: conduit and forwarding-sink are classified", "[taint][interproc]") {
    auto unit = parse(std::string(kDecls) + "char* wrap(char* s){ return s; }\n"
                                            "void run(char* s){ system(s); }\n");
    REQUIRE(unit);
    auto summaries = fold::detect::computeTaintSummaries(unit->getASTContext());

    const auto* wrap = findFn(*unit, "wrap");
    const auto* run = findFn(*unit, "run");
    REQUIRE(wrap);
    REQUIRE(run);
    REQUIRE(summaries.at(wrap).return_tainted_from.count(0) == 1);
    REQUIRE(summaries.at(run).sink_params.count(0) == 1);
    REQUIRE((summaries.at(run).sink_params.at(0) & fold::detect::taint_domain::kCommand) != 0);
}

TEST_CASE("interproc: taint through a conduit reaches a sink", "[taint][interproc]") {
    auto unit = parse(std::string(kDecls) +
                      "char* wrap(char* s){ return s; }\n"
                      "void f(){ char* u = getenv(\"X\"); char* q = wrap(u); system(q); }\n");
    REQUIRE(unit);
    REQUIRE(runTaint(*unit) >= 1);
}

TEST_CASE("interproc: taint into a forwarding sink fires", "[taint][interproc]") {
    auto unit = parse(std::string(kDecls) + "void run(char* s){ system(s); }\n"
                                            "void f(){ char* u = getenv(\"X\"); run(u); }\n");
    REQUIRE(unit);
    REQUIRE(runTaint(*unit) >= 1);
}

TEST_CASE("interproc: a helper that sanitizes before its sink does not fire",
          "[taint][interproc]") {
    auto unit = parse(std::string(kDecls) +
                      "void runSql(MYSQL* db, char* s){ char* c = PQescapeLiteral(0, s, 0);"
                      " mysql_query(db, c); }\n"
                      "void f(MYSQL* db){ char* u = getenv(\"X\"); runSql(db, u); }\n");
    REQUIRE(unit);
    // runSql escapes `s` for SQL before the SQL sink, so its sink summary is
    // empty and the forwarding call produces no finding.
    REQUIRE(runTaint(*unit) == 0);
}

TEST_CASE("interproc: a non-forwarding helper does not taint", "[taint][interproc]") {
    auto unit = parse(std::string(kDecls) +
                      "int len(char* s){ int n = 0; while (s[n]) n++; return n; }\n"
                      "void f(){ char* u = getenv(\"X\"); int n = len(u); (void)n; }\n");
    REQUIRE(unit);
    REQUIRE(runTaint(*unit) == 0);
}