#include <clang/AST/Decl.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/symbolic/encoder.hpp"

namespace {

void requireConcreteReturn(const std::string& code, std::uint64_t expected) {
    auto unit = clang::tooling::buildASTFromCodeWithArgs(code, {"-std=c++20"}, "input.cpp");
    REQUIRE(unit);

    auto& ctx = unit->getASTContext();
    const clang::FunctionDecl* target = nullptr;
    for (auto* decl : ctx.getTranslationUnitDecl()->decls()) {
        const auto* fd = clang::dyn_cast<clang::FunctionDecl>(decl);
        if (fd && fd->hasBody() && fd->getNameAsString() == "f") {
            target = fd;
            break;
        }
    }
    REQUIRE(target != nullptr);

    auto env = fold::SymbolicEnv::create();
    const auto* body = clang::dyn_cast<clang::CompoundStmt>(target->getBody());
    REQUIRE(body != nullptr);
    for (const clang::Stmt* stmt : body->body()) {
        if (const auto* ret = clang::dyn_cast<clang::ReturnStmt>(stmt)) {
            z3::expr value = fold::encodeExpr(ret->getRetValue(), ctx, env);
            z3::solver solver(value.ctx());
            solver.add(value != value.ctx().bv_val(expected, 64));
            REQUIRE(solver.check() == z3::unsat);
            return;
        }
        (void)fold::applyStmt(stmt, ctx, env);
    }

    FAIL("function body did not return");
}

} // namespace

TEST_CASE("symbolic arrays: memcpy propagates stack bytes", "[symbolic][array][memcpy][stack]") {
    requireConcreteReturn(
        "extern \"C\" void* memcpy(void*, const void*, unsigned long);\n"
        "int f(){ char src[4]; char dst[4]; src[0] = 'A'; memcpy(dst, src, 1); return dst[0]; }",
        static_cast<std::uint64_t>('A'));
}

TEST_CASE("symbolic arrays: strcpy writes literal bytes and terminator",
          "[symbolic][array][strcpy][stack]") {
    requireConcreteReturn("extern \"C\" char* strcpy(char*, const char*);\n"
                          "int f(){ char dst[8]; strcpy(dst, \"ABCD\"); return dst[4]; }",
                          0);
}

TEST_CASE("symbolic arrays: memcpy propagates heap bytes", "[symbolic][array][memcpy][heap]") {
    requireConcreteReturn(
        "extern \"C\" void* malloc(unsigned long);\n"
        "extern \"C\" void* memcpy(void*, const void*, unsigned long);\n"
        "int f(){ char* src = (char*)malloc(4); char* dst = (char*)malloc(4); src[0] = 'Q'; "
        "memcpy(dst, src, 1); return dst[0]; }",
        static_cast<std::uint64_t>('Q'));
}