#include <clang/AST/ASTContext.h>
#include <clang/AST/Decl.h>
#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/detect/detail/security_walker_analysis.hpp"
#include "fold/detect/detail/security_walker_refinement.hpp"
#include "fold/detect/security_walker.hpp"
#include "fold/symbolic/safety_observer.hpp"

namespace {

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

class FunctionFinder final : public clang::RecursiveASTVisitor<FunctionFinder> {
  public:
    explicit FunctionFinder(std::string name) : name_(std::move(name)) {}

    bool VisitFunctionDecl(clang::FunctionDecl* function) {
        if (function && function->hasBody() && function->getNameAsString() == name_)
            function_ = function;
        return true;
    }

    clang::FunctionDecl* function() const { return function_; }

  private:
    std::string name_;
    clang::FunctionDecl* function_ = nullptr;
};

class ForFinder final : public clang::RecursiveASTVisitor<ForFinder> {
  public:
    bool VisitForStmt(clang::ForStmt* loop) {
        if (!loop_)
            loop_ = loop;
        return true;
    }

    clang::ForStmt* loop() const { return loop_; }

  private:
    clang::ForStmt* loop_ = nullptr;
};

clang::FunctionDecl* findFunction(clang::ASTContext& context, const std::string& name) {
    FunctionFinder finder(name);
    finder.TraverseDecl(context.getTranslationUnitDecl());
    return finder.function();
}

clang::ForStmt* findFor(clang::FunctionDecl* function) {
    if (!function)
        return nullptr;
    ForFinder finder;
    finder.TraverseStmt(function->getBody());
    return finder.loop();
}

class ThrowOnceDivisionObserver final : public fold::SafetyObserver {
  public:
    void onFunctionStart(const clang::FunctionDecl*, z3::solver*) override { ++starts; }
    void onFunctionEnd(const clang::FunctionDecl*) override { ++ends; }

    void onDivision(const z3::expr&, const z3::expr&, const fold::ir::Expr&, const fold::ir::Expr&,
                    const z3::expr&, const fold::ir::SourceLoc&, bool, unsigned,
                    const fold::DefUseContext*, const fold::SymbolicEnv*) override {
        if (!threw) {
            threw = true;
            throw z3::exception("synthetic unsupported translator expression");
        }
    }

    unsigned starts = 0;
    unsigned ends = 0;
    bool threw = false;
};

} // namespace

TEST_CASE("security walker isolates symbolic translation failures per function",
          "[security-walker][translator]") {
    auto unit = parse("int first(int x) { return 10 / x; } "
                      "int second(int x) { return 20 / x; }");
    REQUIRE(unit);
    ThrowOnceDivisionObserver observer;
    fold::SolverBudget budget;
    fold::SolverStats stats;

    CHECK_NOTHROW(fold::walkContextWithObserver(unit->getASTContext(), observer, budget, stats));
    CHECK(observer.threw);
    CHECK(observer.starts == 2);
    CHECK(observer.ends == 2);
    CHECK(stats.bailed_functions == 1);
    CHECK(stats.unknown == 1);
}

TEST_CASE("security sink precheck distinguishes sink families", "[security-walker]") {
    auto unit = parse(R"cpp(
        int divide(int value, int divisor) { return value / divisor; }
        int add(int lhs, int rhs) { return lhs + rhs; }
    )cpp");
    REQUIRE(unit);
    auto& context = unit->getASTContext();

    auto* divide = findFunction(context, "divide");
    auto* add = findFunction(context, "add");
    REQUIRE(divide);
    REQUIRE(add);

    const auto division = fold::security_walker_detail::precheckSecuritySinks(
        divide, fold::SecuritySinkKind::DivByZero);
    const auto wrong_family =
        fold::security_walker_detail::precheckSecuritySinks(divide, fold::SecuritySinkKind::OOB);
    const auto clean =
        fold::security_walker_detail::precheckSecuritySinks(add, fold::SecuritySinkKind::DivByZero);

    REQUIRE(division.has_candidate);
    REQUIRE(division.candidate_sink_count == 1);
    REQUIRE_FALSE(wrong_family.has_candidate);
    REQUIRE_FALSE(clean.has_candidate);
    REQUIRE(division.stmt_count > 0);
}

TEST_CASE("security sink precheck recognizes every SQL sink", "[security-walker][sql]") {
    auto unit = parse(R"cpp(
        struct DB { void execute(const char*); void exec(const char*); void query(const char*); };
        void mysql_query(int, const char*);
        void PQexec(int, const char*);
        void PQexecParams(int, const char*);
        void sqlite3_exec(int, const char*);
        void sqlite3_prepare_v2(int, const char*);
        void sqlite3_prepare_v3(int, const char*);
        void sqlite3_prepare(int, const char*);
        void mysql() { mysql_query(0, "q"); }
        void pg() { PQexec(0, "q"); }
        void pg_params() { PQexecParams(0, "q"); }
        void sqlite_exec() { sqlite3_exec(0, "q"); }
        void sqlite_v2() { sqlite3_prepare_v2(0, "q"); }
        void sqlite_v3() { sqlite3_prepare_v3(0, "q"); }
        void sqlite() { sqlite3_prepare(0, "q"); }
        void member_execute(DB& db) { db.execute("q"); }
        void member_exec(DB& db) { db.exec("q"); }
        void member_query(DB& db) { db.query("q"); }
        void unrelated() { int x = 1; (void)x; }
    )cpp");
    REQUIRE(unit);
    auto& context = unit->getASTContext();

    for (const auto* name : {"mysql", "pg", "pg_params", "sqlite_exec", "sqlite_v2", "sqlite_v3",
                             "sqlite", "member_execute", "member_exec", "member_query"}) {
        const auto result = fold::security_walker_detail::precheckSecuritySinks(
            findFunction(context, name), fold::SecuritySinkKind::SqlInjection);
        CHECK(result.has_candidate);
        CHECK(result.candidate_sink_count == 1);
    }
    const auto clean = fold::security_walker_detail::precheckSecuritySinks(
        findFunction(context, "unrelated"), fold::SecuritySinkKind::SqlInjection);
    CHECK_FALSE(clean.has_candidate);
}

TEST_CASE("counted-loop refinement is precise and conservative", "[security-walker]") {
    auto unit = parse(R"cpp(
        void unit_stride() { for (int i = 2; i < 5; ++i) {} }
        void strided() { for (int i = 2; i < 8; i += 2) {} }
    )cpp");
    REQUIRE(unit);
    auto& context = unit->getASTContext();

    const auto precise = fold::security_walker_detail::analyzeForLoopRange(
        findFor(findFunction(context, "unit_stride")), context);
    const auto conservative = fold::security_walker_detail::analyzeForLoopRange(
        findFor(findFunction(context, "strided")), context);

    REQUIRE(precise);
    REQUIRE(precise->range.lo == 2);
    REQUIRE(precise->range.hi == 4);
    REQUIRE_FALSE(precise->range.top);
    REQUIRE_FALSE(conservative);
}

TEST_CASE("interval joins preserve top and disjoint uncertainty", "[security-walker]") {
    const fold::IntInterval low{0, 4, false};
    const fold::IntInterval high{8, 12, false};

    const auto overlap =
        fold::security_walker_detail::intersectIntervals(low, fold::IntInterval{2, 9, false});
    REQUIRE(overlap.lo == 2);
    REQUIRE(overlap.hi == 4);
    REQUIRE_FALSE(overlap.top);

    REQUIRE(fold::security_walker_detail::intersectIntervals(low, high).top);
    REQUIRE(fold::security_walker_detail::unionIntervals(low, high).lo == 0);
    REQUIRE(fold::security_walker_detail::unionIntervals(low, high).hi == 12);
}