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

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

#include "fold/symbolic/loop_inventory.hpp"

namespace {

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

const clang::FunctionDecl* findFunction(clang::ASTUnit& unit, const std::string& name) {
    auto& ctx = unit.getASTContext();
    for (auto* decl : ctx.getTranslationUnitDecl()->decls()) {
        auto* fd = clang::dyn_cast<clang::FunctionDecl>(decl);
        if (!fd || !fd->hasBody())
            continue;
        if (fd->getQualifiedNameAsString() == name)
            return fd;
    }
    return nullptr;
}

fold::symbolic::FunctionCostSummary summarize(const std::string& code,
                                              const std::string& function = "f") {
    auto unit = parse(code);
    REQUIRE(unit);
    auto* fn = findFunction(*unit, function);
    REQUIRE(fn != nullptr);
    return fold::symbolic::collectLoopInventoryForFunction(fn, unit->getASTContext());
}

} // namespace

TEST_CASE("loop inventory reports O(1) functions without loops", "[loop-inventory]") {
    auto summary = summarize("int f(int x) { return x + 1; }");

    REQUIRE(summary.function == "f");
    REQUIRE(summary.backend == "source");
    REQUIRE(summary.location.isValid());
    REQUIRE(summary.cost_expression == "O(1)");
    REQUIRE(summary.confidence == fold::symbolic::CostSummaryConfidence::Syntactic);
    REQUIRE(summary.precision == fold::EncodingPrecision::Exact);
    REQUIRE(summary.security_relevance == fold::symbolic::CostSecurityRelevance::None);
    REQUIRE(summary.security_reasons.empty());
    REQUIRE(summary.loops.empty());
}

TEST_CASE("loop inventory records linear for-loop structure", "[loop-inventory]") {
    auto summary = summarize(R"(
        void work(int);
        void f(int n) {
            for (int i = 0; i < n; ++i) {
                work(i);
            }
        }
    )");

    REQUIRE(summary.cost_expression == "O(n)");
    REQUIRE(summary.input_symbols.size() == 1);
    REQUIRE(summary.input_symbols[0].name == "n");
    REQUIRE(summary.loops.size() == 1);

    const auto& loop = summary.loops[0];
    REQUIRE(loop.kind == "for");
    REQUIRE(loop.nesting_depth == 1);
    REQUIRE(loop.guard_expression == "i < n");
    REQUIRE(loop.induction_variable == "i");
    REQUIRE(loop.bound_kind == fold::symbolic::CostBoundKind::Linear);
    REQUIRE(loop.bound_expression == "max(0, n)");
    REQUIRE(loop.confidence == fold::symbolic::CostSummaryConfidence::Proved);
    REQUIRE(loop.precision == fold::EncodingPrecision::Exact);
    REQUIRE(loop.precision_detail == "z3-proved-affine-trip-count");
    REQUIRE(loop.security_relevance == fold::symbolic::CostSecurityRelevance::InputDependent);
    REQUIRE(loop.security_reasons ==
            std::vector<std::string>{"loop-bound-depends-on-input-symbol"});
    REQUIRE(summary.security_relevance == fold::symbolic::CostSecurityRelevance::InputDependent);
    REQUIRE(summary.security_reasons ==
            std::vector<std::string>{"loop-bound-depends-on-input-symbol"});
    REQUIRE(loop.calls == std::vector<std::string>{"work"});
    REQUIRE(loop.unsupported_constructs.empty());
}

TEST_CASE("loop inventory records nested symbolic loops structurally", "[loop-inventory]") {
    auto summary = summarize(R"(
        void f(int n, int m) {
            for (int i = 0; i < n; ++i) {
                for (int j = 0; j < m; ++j) {}
            }
        }
    )");

    REQUIRE(summary.cost_expression == "O(n^2)");
    REQUIRE(summary.loops.size() == 2);
    REQUIRE(summary.loops[0].nesting_depth == 1);
    REQUIRE(summary.loops[1].nesting_depth == 2);
    REQUIRE(summary.loops[0].input_symbols[0].name == "n");
    REQUIRE(summary.loops[1].input_symbols[0].name == "m");
    REQUIRE(summary.loops[0].security_relevance ==
            fold::symbolic::CostSecurityRelevance::InputDependent);
    REQUIRE(summary.loops[1].security_relevance ==
            fold::symbolic::CostSecurityRelevance::PotentialDos);
    REQUIRE(summary.security_relevance == fold::symbolic::CostSecurityRelevance::PotentialDos);
}

TEST_CASE("loop inventory proves decreasing affine loop bounds", "[loop-inventory]") {
    auto summary = summarize(R"(
        void f(int n) {
            for (int i = 10; i > n; --i) {}
        }
    )");

    REQUIRE(summary.confidence == fold::symbolic::CostSummaryConfidence::Proved);
    REQUIRE(summary.loops.size() == 1);
    REQUIRE(summary.loops[0].bound_kind == fold::symbolic::CostBoundKind::Linear);
    REQUIRE(summary.loops[0].bound_expression == "max(0, 10 - n)");
    REQUIRE(summary.loops[0].confidence == fold::symbolic::CostSummaryConfidence::Proved);
    REQUIRE(summary.loops[0].precision == fold::EncodingPrecision::Exact);
}

TEST_CASE("loop inventory proves coarse upper bounds for non-unit affine strides",
          "[loop-inventory]") {
    auto summary = summarize(R"(
        void f(int n) {
            for (int i = 0; i < n; i += 2) {}
        }
    )");

    REQUIRE(summary.loops.size() == 1);
    REQUIRE(summary.loops[0].bound_kind == fold::symbolic::CostBoundKind::Linear);
    REQUIRE(summary.loops[0].bound_expression == "max(0, n)");
    REQUIRE(summary.loops[0].confidence == fold::symbolic::CostSummaryConfidence::Proved);
    REQUIRE(summary.loops[0].precision == fold::EncodingPrecision::OverApproximated);
    REQUIRE(summary.loops[0].precision_detail == "z3-proved-affine-trip-upper-bound");
}

TEST_CASE("loop inventory does not prove self-referential affine guards", "[loop-inventory]") {
    auto summary = summarize(R"(
        void f(int n) {
            for (int i = 0; i < i + n; ++i) {}
        }
    )");

    REQUIRE(summary.loops.size() == 1);
    REQUIRE(summary.loops[0].bound_kind == fold::symbolic::CostBoundKind::Linear);
    REQUIRE(summary.loops[0].confidence == fold::symbolic::CostSummaryConfidence::Heuristic);
    REQUIRE(summary.loops[0].precision == fold::EncodingPrecision::OverApproximated);
    REQUIRE(summary.loops[0].precision_detail.empty());
}

TEST_CASE("loop inventory preserves unsupported loop evidence without findings",
          "[loop-inventory]") {
    auto summary = summarize(R"(
        int next();
        void f(int n) {
            while (next() < n) {}
        }
    )");

    REQUIRE(summary.cost_expression == "unknown");
    REQUIRE(summary.confidence == fold::symbolic::CostSummaryConfidence::Syntactic);
    REQUIRE(summary.precision == fold::EncodingPrecision::UnsupportedNode);
    REQUIRE(summary.loops.size() == 1);
    REQUIRE(summary.loops[0].kind == "while");
    REQUIRE(summary.loops[0].guard_expression == "next() < n");
    REQUIRE(summary.loops[0].calls == std::vector<std::string>{"next"});
    REQUIRE(summary.loops[0].security_relevance ==
            fold::symbolic::CostSecurityRelevance::LatentDos);
    REQUIRE(summary.security_relevance == fold::symbolic::CostSecurityRelevance::LatentDos);
    REQUIRE_FALSE(summary.loops[0].unsupported_constructs.empty());
}

TEST_CASE("loop inventory raises context relevance for parser/string loops", "[loop-inventory]") {
    auto summary = summarize(R"(
        void decode_token(int);
        void parse_request(int len) {
            for (int i = 0; i < len; ++i) {
                decode_token(i);
            }
        }
    )",
                             "parse_request");

    REQUIRE(summary.security_relevance == fold::symbolic::CostSecurityRelevance::LatentDos);
    REQUIRE(summary.loops.size() == 1);
    REQUIRE(summary.loops[0].security_reasons ==
            std::vector<std::string>{"loop-bound-depends-on-input-symbol",
                                     "input-dependent-loop-in-request-parse-string-context"});

    auto no_call_summary = summarize(R"(
        void parse_request(int len) {
            int sum = 0;
            for (int i = 0; i < len; ++i) {
                sum += i;
            }
        }
    )",
                                     "parse_request");
    REQUIRE(no_call_summary.security_relevance == fold::symbolic::CostSecurityRelevance::LatentDos);
    REQUIRE(no_call_summary.loops[0].security_reasons ==
            std::vector<std::string>{"loop-bound-depends-on-input-symbol",
                                     "input-dependent-loop-in-request-parse-string-context"});
}

TEST_CASE("loop inventory collector has a bounded microbenchmark", "[loop-inventory][benchmark]") {
    std::ostringstream code;
    code << "void f(int n) {\n";
    for (int idx = 0; idx < 48; ++idx) {
        code << "  for (int i" << idx << " = 0; i" << idx << " < n; ++i" << idx << ") {}\n";
    }
    code << "}\n";

    auto unit = parse(code.str());
    REQUIRE(unit);
    auto* fn = findFunction(*unit, "f");
    REQUIRE(fn != nullptr);

    const auto started = std::chrono::steady_clock::now();
    auto summary = fold::symbolic::collectLoopInventoryForFunction(fn, unit->getASTContext());
    const auto elapsed = std::chrono::steady_clock::now() - started;

    REQUIRE(summary.loops.size() == 48);
    REQUIRE(summary.confidence == fold::symbolic::CostSummaryConfidence::Proved);
    REQUIRE(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count() < 2000);
}

TEST_CASE("loop inventory collects all function summaries in a translation unit",
          "[loop-inventory]") {
    auto unit = parse("int f(){ return 0; } int g(int n){ for(int i=0; i<n; ++i){} return n; }");
    REQUIRE(unit);

    auto summaries = fold::symbolic::collectLoopInventory(unit->getASTContext());
    REQUIRE(summaries.size() == 2);
    REQUIRE(summaries[0].function == "f");
    REQUIRE(summaries[1].function == "g");
    REQUIRE(summaries[0].cost_expression == "O(1)");
    REQUIRE(summaries[1].cost_expression == "O(n)");
}