#include <clang/AST/ASTContext.h>
#include <clang/AST/Decl.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/ir/cfg_walker.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* functionNamed(clang::ASTUnit& unit, const std::string& name) {
    for (const clang::Decl* decl : unit.getASTContext().getTranslationUnitDecl()->decls()) {
        const auto* fd = clang::dyn_cast<clang::FunctionDecl>(decl);
        if (fd && fd->hasBody() && fd->getQualifiedNameAsString() == name)
            return fd;
    }
    return nullptr;
}

} // namespace

TEST_CASE("cfg-walker: builds reachable blocks and edges", "[cfg]") {
    auto unit = parse(R"cpp(
        int f(int x) {
            if (x > 0)
                return x;
            return -x;
        }
    )cpp");
    REQUIRE(unit);

    const clang::FunctionDecl* fd = functionNamed(*unit, "f");
    REQUIRE(fd);

    auto maybe_cfg = fold::ir::buildFunctionCFG(fd, unit->getASTContext());
    REQUIRE(maybe_cfg.has_value());
    REQUIRE(maybe_cfg->cfg != nullptr);
    REQUIRE(maybe_cfg->function_name == "f");

    fold::ir::CFGWalker walker(*maybe_cfg);
    std::vector<unsigned> visited_blocks;
    std::vector<std::pair<unsigned, unsigned>> visited_edges;

    REQUIRE(walker.walk(
        [&](unsigned block_id, const clang::CFGBlock* block) {
            REQUIRE(block != nullptr);
            visited_blocks.push_back(block_id);
            return true;
        },
        [&](unsigned from_id, unsigned to_id) { visited_edges.emplace_back(from_id, to_id); }));

    REQUIRE(!visited_blocks.empty());
    REQUIRE(visited_blocks.front() == maybe_cfg->cfg->getEntry().getBlockID());
    REQUIRE(visited_edges.size() >= 2);
}

TEST_CASE("cfg-walker: walk can be interrupted", "[cfg]") {
    auto unit = parse("int g(int x) { while (x > 0) --x; return x; }");
    REQUIRE(unit);

    const clang::FunctionDecl* fd = functionNamed(*unit, "g");
    REQUIRE(fd);

    auto maybe_cfg = fold::ir::buildFunctionCFG(fd, unit->getASTContext());
    REQUIRE(maybe_cfg.has_value());

    fold::ir::CFGWalker walker(*maybe_cfg);
    unsigned calls = 0;

    REQUIRE_FALSE(walker.walk([&](unsigned, const clang::CFGBlock*) {
        ++calls;
        return false;
    }));
    REQUIRE(calls == 1);
}