#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/cross_tu_cache.hpp"
#include "fold/detect/double_free.hpp"
#include "fold/detect/taint_track.hpp"
#include "fold/detect/use_after_free.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"
#include "fold/symbolic/encoder_callee_names.hpp"

// Library models (malloc/free, taint sources and sinks) were selected by the
// callee's bare identifier, so a project function that merely shares a libc
// name was modeled as libc: `mine::free(p)` followed by `free(p)` was a Proved
// double-free. Models now require the resolved declaration to be a C library
// function (global scope, namespace std, or extern "C"; never a member).

namespace {

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

std::size_t countKind(const std::vector<fold::Finding>& findings, fold::FindingKind kind) {
    std::size_t n = 0;
    for (const auto& f : findings)
        if (f.kind == kind)
            ++n;
    return n;
}

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

std::vector<fold::Finding> runUseAfterFree(const std::string& code) {
    auto unit = parse(code);
    REQUIRE(unit);
    return fold::detectUseAfterFreeInContext(unit->getASTContext());
}

std::vector<fold::Finding> runTaint(const std::string& code) {
    fold::detect::CrossTUCache::instance().clear();
    auto unit = parse(code);
    REQUIRE(unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    return fold::detectTaintTrackInContext(unit->getASTContext(), b, s);
}

const std::string kLibc = "extern \"C\" void* malloc(unsigned long);\n"
                          "extern \"C\" void free(void*);\n";

} // namespace

TEST_CASE("heap model: a namespaced function named free is not libc free",
          "[library-callee][double-free]") {
    const std::string code = kLibc + "namespace mine { inline void free(void*) {} }\n"
                                     "int f() {\n"
                                     "  int* p = static_cast<int*>(malloc(sizeof(int)));\n"
                                     "  if (!p) return 0;\n"
                                     "  *p = 1;\n"
                                     "  mine::free(p);\n"
                                     "  int v = *p;\n"
                                     "  free(p);\n"
                                     "  return v;\n"
                                     "}\n";
    CHECK(countKind(runDoubleFree(code), fold::FindingKind::DoubleFree) == 0);
    CHECK(countKind(runUseAfterFree(code), fold::FindingKind::UseAfterFree) == 0);
}

TEST_CASE("heap model: a member function named free is not libc free",
          "[library-callee][double-free]") {
    const std::string code = kLibc + "struct Pool { void free(void*) {} };\n"
                                     "void f(Pool& pool) {\n"
                                     "  void* p = malloc(8);\n"
                                     "  if (!p) return;\n"
                                     "  pool.free(p);\n"
                                     "  free(p);\n"
                                     "}\n";
    CHECK(countKind(runDoubleFree(code), fold::FindingKind::DoubleFree) == 0);
}

TEST_CASE("heap model: libc free twice is still a double free", "[library-callee][double-free]") {
    // Negative control: the fix must not drop the real model.
    const std::string code = kLibc + "void f() {\n"
                                     "  void* p = malloc(8);\n"
                                     "  if (!p) return;\n"
                                     "  free(p);\n"
                                     "  free(p);\n"
                                     "}\n";
    CHECK(countKind(runDoubleFree(code), fold::FindingKind::DoubleFree) == 1);
}

TEST_CASE("taint: a member function named system is not a command sink",
          "[library-callee][taint]") {
    const std::string code = "extern \"C\" char* getenv(const char*);\n"
                             "struct Shell { int system(const char*) { return 0; } };\n"
                             "int f(Shell& sh) { return sh.system(getenv(\"CMD\")); }\n";
    CHECK(countKind(runTaint(code), fold::FindingKind::TaintTrack) == 0);
}

TEST_CASE("taint: libc system with tainted input is still a sink", "[library-callee][taint]") {
    const std::string code = "extern \"C\" char* getenv(const char*);\n"
                             "extern \"C\" int system(const char*);\n"
                             "int f() { return system(getenv(\"CMD\")); }\n";
    CHECK(countKind(runTaint(code), fold::FindingKind::TaintTrack) == 1);
}

TEST_CASE("heap model: a global C++ function named free is not libc free",
          "[library-callee][double-free][global]") {
    // A global-scope C++ overload (no extern "C", not in a system header) is a
    // project function that shares libc's name, not the library function.
    const std::string code = kLibc + "struct Obj { int v; };\n"
                                     "void free(Obj*) {}\n"
                                     "int f() {\n"
                                     "  Obj* o = static_cast<Obj*>(malloc(sizeof(Obj)));\n"
                                     "  if (!o) return 0;\n"
                                     "  o->v = 1;\n"
                                     "  free(o);\n"
                                     "  int v = o->v;\n"
                                     "  free(o);\n"
                                     "  return v;\n"
                                     "}\n";
    CHECK(countKind(runDoubleFree(code), fold::FindingKind::DoubleFree) == 0);
    CHECK(countKind(runUseAfterFree(code), fold::FindingKind::UseAfterFree) == 0);
}