#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/always_condition.hpp"
#include "fold/detect/dead_branch.hpp"
#include "fold/detect/loop_havoc.hpp"
#include "fold/detect/unreachable.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"
#include "fold/symbolic/modeling_options.hpp"

// Bounded loop unrolling followed only the path on which the loop guard stays
// true. When the guard could also be false (a symbolic bound), that early exit
// was dropped, so the state after the loop was "the loop ran loop_unroll
// times" and conditions on it were reported as Proved dead / always-true or
// always-false. These cases were all Proved before the fix.

namespace {

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

std::vector<fold::Finding> runDead(const std::string& code) {
    auto unit = parse(code);
    REQUIRE(unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    return fold::detectDeadBranchesInContext(unit->getASTContext(), b, s);
}

std::vector<fold::Finding> runAlways(const std::string& code) {
    auto unit = parse(code);
    REQUIRE(unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    return fold::detectAlwaysConditionsInContext(unit->getASTContext(), b, s);
}

std::vector<fold::Finding> runUnreachable(const std::string& code) {
    auto unit = parse(code);
    REQUIRE(unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    return fold::detectUnreachableCodeInContext(unit->getASTContext(), b, s);
}

std::size_t proved(const std::vector<fold::Finding>& findings) {
    std::size_t n = 0;
    for (const auto& f : findings)
        if (f.confidence == fold::Confidence::Proved)
            ++n;
    return n;
}

// Symbolic bound: the loop may run zero times, or more than loop_unroll times.
const std::string kSymbolicLoops =
    "int a0(unsigned n) { int s = 0; for (unsigned i = 0; i < n; ++i) { s += 1; }"
    "  if (s == 0) return 7; return 0; }\n"
    "int a1(unsigned n) { int s = 0; for (unsigned i = 0; i < n; ++i) { s += 1; }"
    "  if (s > 100) return 7; return 0; }\n"
    "int a2(unsigned n) { int s = 0; while (n > 0) { s += 1; n--; }"
    "  if (s == 0) return 7; return 0; }\n"
    "int a3(unsigned n) { int s = 0; for (unsigned i = 0; i < n; ++i) { s += 1; }"
    "  if (s > 5) return 1; return 0; }\n";

// Constant bound within the unroll budget: the unroll sees the loop's only
// exit, so the post-loop state is exact and a genuinely dead branch stays
// provable.
const std::string kExactLoop = "int e(int x) { int s = 0; for (int i = 0; i < 3; ++i) { s += 1; }"
                               "  if (s == 7) return -1; return s + x; }\n";

// Constant bound, but a `break` offers a second way out whose state the
// unroll does not follow.
const std::string kBreakLoop =
    "int b(int x) { int s = 0; for (int i = 0; i < 3; ++i) { if (x == i) break; s += 1; }"
    "  if (s == 0) return -1; return s; }\n";

} // namespace

TEST_CASE("dead: a branch after a symbolic-bound loop is not proven dead", "[dead][loop]") {
    CHECK(proved(runDead(kSymbolicLoops)) == 0);
}

TEST_CASE("always-condition: a condition after a symbolic-bound loop is not proven constant",
          "[always][loop]") {
    CHECK(proved(runAlways(kSymbolicLoops)) == 0);
}

TEST_CASE("unreachable: code after a symbolic-bound loop is not proven unreachable",
          "[unreachable][loop]") {
    CHECK(proved(runUnreachable(kSymbolicLoops)) == 0);
}

TEST_CASE("dead: a loop the unroll exits exactly keeps its precision", "[dead][loop]") {
    // Negative control for the fix: havocking every loop would silence this too.
    CHECK(proved(runDead(kExactLoop)) == 1);
}

TEST_CASE("dead: a break out of a constant loop is an exit the unroll did not follow",
          "[dead][loop]") {
    CHECK(proved(runDead(kBreakLoop)) == 0);
}

namespace {

// The body of the first loop statement in the translation unit.
const clang::Stmt* firstLoopBody(clang::ASTUnit& unit) {
    struct Finder : clang::RecursiveASTVisitor<Finder> {
        const clang::Stmt* body = nullptr;
        bool VisitForStmt(clang::ForStmt* s) {
            if (!body)
                body = s->getBody();
            return true;
        }
        bool VisitWhileStmt(clang::WhileStmt* s) {
            if (!body)
                body = s->getBody();
            return true;
        }
    } finder;
    finder.TraverseDecl(unit.getASTContext().getTranslationUnitDecl());
    return finder.body;
}

bool escapesIn(const std::string& code) {
    auto unit = parse(code);
    REQUIRE(unit);
    const auto* body = firstLoopBody(*unit);
    REQUIRE(body);
    return fold::detect::containsLoopEscape(body);
}

} // namespace

TEST_CASE("loop escape: a break out of the loop escapes; a nested one does not", "[loop_havoc]") {
    CHECK(escapesIn("void f(int x){ for(int i=0;i<3;++i){ if(x) break; } }"));
    CHECK(escapesIn("void f(int x){ while(x){ goto out; } out:; }"));
    CHECK_FALSE(escapesIn("void f(int x){ for(int i=0;i<3;++i){ for(;;){ break; } } }"));
    CHECK_FALSE(escapesIn(
        "void f(int x){ for(int i=0;i<3;++i){ switch(x){ case 1: break; default: break; } } }"));
    CHECK_FALSE(escapesIn("void f(int x){ for(int i=0;i<3;++i){ x += i; } }"));
}

TEST_CASE("loop exits: exact only when the unroll leaves along the only exit", "[loop_havoc]") {
    SECTION("guard falsified after iterations that could not exit") {
        fold::detect::LoopUnrollExits exits(nullptr);
        CHECK(exits.guard(/*may_continue=*/true, /*may_exit=*/false));
        CHECK_FALSE(exits.guard(/*may_continue=*/false, /*may_exit=*/true));
        CHECK(exits.exact());
    }
    SECTION("an iteration that could also exit makes the state inexact") {
        fold::detect::LoopUnrollExits exits(nullptr);
        CHECK(exits.guard(true, true));
        CHECK_FALSE(exits.guard(false, true));
        CHECK_FALSE(exits.exact());
    }
    SECTION("the unroll budget running out is not an exit") {
        fold::detect::LoopUnrollExits exits(nullptr);
        CHECK(exits.guard(true, false));
        CHECK_FALSE(exits.exact());
    }
    SECTION("an unmodeled guard or a stopped body is never exact") {
        fold::detect::LoopUnrollExits unmodeled(nullptr);
        unmodeled.guardNotModeled();
        CHECK_FALSE(unmodeled.guard(false, true));
        CHECK_FALSE(unmodeled.exact());
        fold::detect::LoopUnrollExits stopped(nullptr);
        stopped.bodyStopped();
        CHECK_FALSE(stopped.guard(false, true));
        CHECK_FALSE(stopped.exact());
    }
}

// Review finding: a `continue` arm is dropped by the walkers (the arm always
// exits the iteration), so iterations that took it are not modeled even when
// the guard sequence is exact. `x == 2` holds only when y == 0.
const std::string kContinueLoop =
    "int c(int y) { int x = 0; for (int i = 0; i < 2; i++) { if (y) continue; x++; }"
    "  if (x == 2) return 1; return 0; }\n";

TEST_CASE("a loop with continue is not treated as an exact unroll", "[loop][continue]") {
    CHECK(proved(runAlways(kContinueLoop)) == 0);
    CHECK(proved(runDead(kContinueLoop)) == 0);
    CHECK(proved(runUnreachable(kContinueLoop)) == 0);
}

// Review finding: a truncated unroll havocs scalar variables but kept array
// element writes from the partial unroll. With n == 0, buf[0] is never set.
const std::string kArrayWriteLoop = "int w(unsigned n) { int buf[4] = {0, 0, 0, 0};"
                                    "  for (unsigned i = 0; i < n && i < 4; i++) buf[i] = 1;"
                                    "  if (buf[0] == 1) return 1; return 0; }\n";

// The same through a pointer: the store may have reached any heap array.
const std::string kPointerWriteLoop =
    "#include <cstdlib>\n"
    "int pw(unsigned n) { int* buf = static_cast<int*>(std::calloc(4, sizeof(int)));"
    "  if (!buf) return 0;"
    "  for (unsigned i = 0; i < n && i < 4; i++) buf[i] = 1;"
    "  int r = buf[0] == 1 ? 1 : 0; std::free(buf); return r; }\n";

TEST_CASE("array writes in a truncated loop do not survive into proofs", "[loop][memory]") {
    CHECK(proved(runAlways(kArrayWriteLoop)) == 0);
    CHECK(proved(runDead(kArrayWriteLoop)) == 0);
    CHECK(proved(runUnreachable(kArrayWriteLoop)) == 0);
    CHECK(proved(runAlways(kPointerWriteLoop)) == 0);
    CHECK(proved(runDead(kPointerWriteLoop)) == 0);
}

// Review finding: unreachable omitted the opaque-loop rule. A call can change
// the loop variable through a pointer, which assigned-variable collection
// does not see.
const std::string kOpaqueCallLoop = "void advance(int* p);\n"
                                    "int u() { int i = 0; while (i < 3) { advance(&i); i++; }"
                                    "  if (i != 3) { int k = i * 2; return k; } return 0; }\n";

TEST_CASE("unreachable havocs loops that contain calls", "[unreachable][loop][opaque]") {
    CHECK(proved(runUnreachable(kOpaqueCallLoop)) == 0);
    CHECK(proved(runDead(kOpaqueCallLoop)) == 0);
}