#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/ir/loop.hpp"

namespace {

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

// Grab the first loop statement (For/While/Do) in the TU.
const clang::Stmt* firstLoop(clang::ASTContext& ctx) {
    struct V : clang::RecursiveASTVisitor<V> {
        const clang::Stmt* found = nullptr;
        bool VisitForStmt(clang::ForStmt* s) { return keep(s); }
        bool VisitWhileStmt(clang::WhileStmt* s) { return keep(s); }
        bool VisitDoStmt(clang::DoStmt* s) { return keep(s); }
        bool keep(const clang::Stmt* s) {
            if (!found)
                found = s;
            return found == nullptr;
        }
    } v;
    v.TraverseDecl(ctx.getTranslationUnitDecl());
    return v.found;
}

fold::ir::Loop loopOf(clang::ASTUnit& unit) {
    const clang::Stmt* s = firstLoop(unit.getASTContext());
    REQUIRE(s != nullptr);
    return fold::ir::buildLoopFromClang(s, unit.getASTContext());
}

} // namespace

TEST_CASE("ir::Loop: counted ascending for is summarizable with known trip", "[ir][loop]") {
    auto u = parse("void f(){ for (int i = 0; i < 10; ++i) { (void)i; } }");
    REQUIRE(u);
    auto loop = loopOf(*u);
    REQUIRE(loop.kind == fold::ir::LoopKind::Counted);
    REQUIRE(loop.isSummarizable());
    REQUIRE(loop.trip_count.has_value());
    REQUIRE(*loop.trip_count == 10);
    REQUIRE(loop.induction_vars.size() == 1);
    REQUIRE(loop.induction_vars[0].init == 0);
    REQUIRE(loop.induction_vars[0].stride == 1);
}

TEST_CASE("ir::Loop: stride and <= bound are accounted for", "[ir][loop]") {
    auto u = parse("void f(){ for (int i = 2; i <= 10; i += 2) { (void)i; } }");
    REQUIRE(u);
    auto loop = loopOf(*u);
    REQUIRE(loop.kind == fold::ir::LoopKind::Counted);
    // i = 2,4,6,8,10 -> 5 iterations
    REQUIRE(*loop.trip_count == 5);
    REQUIRE(loop.induction_vars[0].stride == 2);
}

TEST_CASE("ir::Loop: descending counted loop", "[ir][loop]") {
    auto u = parse("void f(){ for (int i = 7; i > 0; --i) { (void)i; } }");
    REQUIRE(u);
    auto loop = loopOf(*u);
    REQUIRE(loop.kind == fold::ir::LoopKind::Counted);
    REQUIRE(*loop.trip_count == 7);
    REQUIRE(loop.induction_vars[0].stride == -1);
}

TEST_CASE("ir::Loop: while(1) scanner is General (fallback)", "[ir][loop]") {
    auto u = parse("void f(unsigned char* p){ while (1) { if (*p == 0) break; ++p; } }");
    REQUIRE(u);
    auto loop = loopOf(*u);
    REQUIRE(loop.kind == fold::ir::LoopKind::General);
    REQUIRE_FALSE(loop.isSummarizable());
}

TEST_CASE("ir::Loop: data-dependent bound is Affine (symbolic trip)", "[ir][loop]") {
    auto u = parse("void f(int n){ for (int i = 0; i < n; ++i) { (void)i; } }");
    REQUIRE(u);
    auto loop = loopOf(*u);
    REQUIRE(loop.kind == fold::ir::LoopKind::Affine);
    REQUIRE_FALSE(loop.trip_count.has_value());
    REQUIRE(loop.induction_vars.size() == 1);
    REQUIRE(loop.induction_vars[0].init == 0);
    REQUIRE(loop.induction_vars[0].stride == 1);
    REQUIRE(loop.isSummarizable());
}