// HNSW Entity Search Benchmark
// Tests HNSW search performance with entity vectors (symbol embeddings)
// Validates search quality and latency for semantic symbol search

#include <catch2/catch_test_macros.hpp>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <map>
#include <numeric>
#include <random>
#include <set>
#include <string>
#include <vector>

#include <nlohmann/json.hpp>

#include <yams/compat/dlfcn.h>
#include <yams/plugins/abi.h>
#include <yams/plugins/symbol_extractor_v1.h>
#include <yams/vector/vector_database.h>

namespace fs = std::filesystem;
using json = nlohmann::json;
using namespace yams::vector;

namespace {

// ============================================================================
// Benchmark Results
// ============================================================================

struct HNSWBenchmarkResult {
    std::string name;
    size_t corpus_size = 0;
    size_t dimension = 0;
    size_t k = 0;
    double mean_latency_us = 0.0;
    double p50_latency_us = 0.0;
    double p95_latency_us = 0.0;
    double p99_latency_us = 0.0;
    double qps = 0.0;
    double recall = 0.0;
};

struct EntitySearchResult {
    std::string name;
    size_t total_entities = 0;
    size_t total_files = 0;
    double extraction_time_ms = 0.0;
    double indexing_time_ms = 0.0;
    double search_latency_us = 0.0;
    double qps = 0.0;
};

// ============================================================================
// Vector Generation Helpers
// ============================================================================

std::vector<float> generateRandomVector(size_t dim, std::mt19937& rng, bool normalize = true) {
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    std::vector<float> vec(dim);
    for (auto& v : vec) {
        v = dist(rng);
    }
    if (normalize) {
        float norm = 0.0f;
        for (float v : vec) {
            norm += v * v;
        }
        norm = std::sqrt(norm);
        if (norm > 0) {
            for (float& v : vec) {
                v /= norm;
            }
        }
    }
    return vec;
}

std::vector<std::vector<float>> generateRandomVectors(size_t count, size_t dim, std::mt19937& rng) {
    std::vector<std::vector<float>> vectors;
    vectors.reserve(count);
    for (size_t i = 0; i < count; ++i) {
        vectors.push_back(generateRandomVector(dim, rng));
    }
    return vectors;
}

// Compute exact k-NN using brute force with cosine similarity (ground truth)
// For normalized vectors, cosine similarity = dot product
std::vector<size_t> bruteForceKNN(const std::vector<float>& query,
                                  const std::vector<std::vector<float>>& corpus, size_t k) {
    std::vector<std::pair<size_t, float>> similarities;
    similarities.reserve(corpus.size());

    for (size_t i = 0; i < corpus.size(); ++i) {
        float dot = 0.0f;
        for (size_t j = 0; j < query.size(); ++j) {
            dot += query[j] * corpus[i][j];
        }
        similarities.emplace_back(i, dot);
    }

    // Sort by descending similarity (highest first)
    std::partial_sort(similarities.begin(), similarities.begin() + std::min(k, similarities.size()),
                      similarities.end(),
                      [](const auto& a, const auto& b) { return a.second > b.second; });

    std::vector<size_t> result;
    result.reserve(k);
    for (size_t i = 0; i < std::min(k, similarities.size()); ++i) {
        result.push_back(similarities[i].first);
    }
    return result;
}

// Calculate recall@k
double calculateRecall(const std::vector<std::string>& retrieved,
                       const std::vector<size_t>& ground_truth_indices,
                       const std::vector<std::string>& all_chunk_ids) {
    std::set<std::string> gt_set;
    for (size_t idx : ground_truth_indices) {
        if (idx < all_chunk_ids.size()) {
            gt_set.insert(all_chunk_ids[idx]);
        }
    }

    size_t hits = 0;
    for (const auto& id : retrieved) {
        if (gt_set.count(id)) {
            hits++;
        }
    }

    return gt_set.empty() ? 0.0 : static_cast<double>(hits) / gt_set.size();
}

// ============================================================================
// Symbol Extraction Plugin
// ============================================================================

struct PluginHandle {
    void* handle = nullptr;
    yams_symbol_extractor_v1* api = nullptr;

    ~PluginHandle() {
        // Don't close handle in benchmark
    }
};

std::optional<PluginHandle> loadSymbolExtractorPlugin() {
    PluginHandle p;

#ifdef __APPLE__
    const char* libname = "yams_symbol_extractor.dylib";
#else
    const char* libname = "yams_symbol_extractor.so";
#endif

    std::vector<std::string> paths;
    const char* buildroot = std::getenv("MESON_BUILD_ROOT");
    if (buildroot && *buildroot) {
        paths.push_back(std::string(buildroot) + "/plugins/symbol_extractor_treesitter/" + libname);
    }
    paths.push_back("plugins/symbol_extractor_treesitter/" + std::string(libname));
    paths.push_back("./plugins/symbol_extractor_treesitter/" + std::string(libname));
    paths.push_back("../plugins/symbol_extractor_treesitter/" + std::string(libname));
    paths.push_back("../../plugins/symbol_extractor_treesitter/" + std::string(libname));

    for (const auto& path : paths) {
        p.handle = dlopen(path.c_str(), RTLD_LAZY | RTLD_LOCAL);
        if (!p.handle) {
            continue;
        }

        auto get_iface = reinterpret_cast<int (*)(const char*, uint32_t, void**)>(
            dlsym(p.handle, "yams_plugin_get_interface"));
        if (!get_iface) {
            continue;
        }

        void* iface_ptr = nullptr;
        int rc = get_iface(YAMS_IFACE_SYMBOL_EXTRACTOR_V1, YAMS_IFACE_SYMBOL_EXTRACTOR_V1_VERSION,
                           &iface_ptr);
        if (rc != YAMS_PLUGIN_OK || !iface_ptr) {
            continue;
        }

        p.api = static_cast<yams_symbol_extractor_v1*>(iface_ptr);

        auto init_fn = reinterpret_cast<int (*)(const char*, const void*)>(
            dlsym(p.handle, "yams_plugin_init"));
        if (init_fn) {
            init_fn(nullptr, nullptr);
        }

        return p;
    }

    return std::nullopt;
}

// Extract symbols from a C++ file
std::vector<std::pair<std::string, std::string>> extractSymbols(PluginHandle& plugin,
                                                                const fs::path& file_path) {
    std::vector<std::pair<std::string, std::string>> symbols; // {name, qualified_name}

    if (!fs::exists(file_path) || !fs::is_regular_file(file_path)) {
        return symbols;
    }

    std::ifstream file(file_path, std::ios::binary);
    std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    if (content.empty()) {
        return symbols;
    }

    yams_symbol_extraction_result_v1* result = nullptr;
    int rc = plugin.api->extract_symbols(plugin.api->self, content.c_str(), content.size(),
                                         file_path.string().c_str(), "cpp", &result);

    if (rc != YAMS_PLUGIN_OK || !result) {
        return symbols;
    }

    for (size_t i = 0; i < result->symbol_count; ++i) {
        std::string name = result->symbols[i].name ? result->symbols[i].name : "";
        std::string qname =
            result->symbols[i].qualified_name ? result->symbols[i].qualified_name : name;
        if (!name.empty()) {
            symbols.emplace_back(name, qname);
        }
    }

    plugin.api->free_result(plugin.api->self, result);
    return symbols;
}

// Collect C++ source files
std::vector<fs::path> collectSourceFiles(const fs::path& dir, size_t max_files = 100) {
    std::vector<fs::path> files;
    if (!fs::exists(dir)) {
        return files;
    }

    for (const auto& entry :
         fs::recursive_directory_iterator(dir, fs::directory_options::skip_permission_denied)) {
        if (!entry.is_regular_file()) {
            continue;
        }
        auto ext = entry.path().extension().string();
        if (ext == ".cpp" || ext == ".hpp" || ext == ".h" || ext == ".cc") {
            files.push_back(entry.path());
            if (files.size() >= max_files) {
                break;
            }
        }
    }
    return files;
}

// Write benchmark results to JSON
void writeBenchmarkJSON(const fs::path& output, const std::vector<HNSWBenchmarkResult>& results) {
    json arr = json::array();
    for (const auto& r : results) {
        arr.push_back({{"name", r.name},
                       {"corpus_size", r.corpus_size},
                       {"dimension", r.dimension},
                       {"k", r.k},
                       {"mean_latency_us", r.mean_latency_us},
                       {"p50_latency_us", r.p50_latency_us},
                       {"p95_latency_us", r.p95_latency_us},
                       {"p99_latency_us", r.p99_latency_us},
                       {"qps", r.qps},
                       {"recall", r.recall}});
    }

    fs::create_directories(output.parent_path());
    std::ofstream out(output);
    out << arr.dump(2);
}

void writeEntitySearchJSON(const fs::path& output, const EntitySearchResult& result) {
    json j = {{"name", result.name},
              {"total_entities", result.total_entities},
              {"total_files", result.total_files},
              {"extraction_time_ms", result.extraction_time_ms},
              {"indexing_time_ms", result.indexing_time_ms},
              {"search_latency_us", result.search_latency_us},
              {"qps", result.qps},
              {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}};

    fs::create_directories(output.parent_path());
    std::ofstream out(output);
    out << j.dump(2);
}

// ============================================================================
// Benchmark Fixtures
// ============================================================================

struct HNSWSearchFixture {
    VectorDatabaseConfig createConfig(size_t dim) {
        VectorDatabaseConfig config;
        config.database_path = ":memory:";
        config.embedding_dim = dim;
        config.create_if_missing = true;
        config.use_in_memory = true;
        return config;
    }

    VectorRecord createRecord(const std::string& chunk_id, const std::string& doc_hash,
                              const std::vector<float>& embedding) {
        VectorRecord rec;
        rec.chunk_id = chunk_id;
        rec.document_hash = doc_hash;
        rec.embedding = embedding;
        rec.content = "Test content for " + chunk_id;
        rec.start_offset = 0;
        rec.end_offset = rec.content.size();
        return rec;
    }
};

} // namespace

// ============================================================================
// HNSW Search Speed Benchmarks
// ============================================================================

TEST_CASE_METHOD(HNSWSearchFixture, "HNSW Search Speed Benchmark",
                 "[benchmark][hnsw][search][!benchmark]") {
    std::vector<HNSWBenchmarkResult> all_results;
    std::mt19937 rng(42);

    SECTION("Search Latency vs Corpus Size") {
        // Larger corpus sizes for meaningful benchmarks
        std::vector<size_t> corpus_sizes = {1000, 5000, 10000, 50000};
        size_t dim = 384;
        size_t k = 10;

        for (size_t corpus_size : corpus_sizes) {
            INFO("Testing corpus size: " << corpus_size);

            auto config = createConfig(dim);
            VectorDatabase db(config);
            REQUIRE(db.initialize());

            // Generate and insert vectors
            auto vectors = generateRandomVectors(corpus_size, dim, rng);
            std::vector<std::string> chunk_ids;
            chunk_ids.reserve(corpus_size);

            for (size_t i = 0; i < corpus_size; ++i) {
                std::string chunk_id = "chunk_" + std::to_string(i);
                chunk_ids.push_back(chunk_id);
                db.insertVector(createRecord(chunk_id, "doc_" + std::to_string(i), vectors[i]));
            }

            // Generate query and compute ground truth
            auto query = generateRandomVector(dim, rng);
            auto gt_indices = bruteForceKNN(query, vectors, k);

            // Benchmark search
            std::vector<double> latencies_us;
            size_t num_iterations = 100;

            for (size_t iter = 0; iter < num_iterations; ++iter) {
                auto start = std::chrono::high_resolution_clock::now();

                VectorSearchParams params;
                params.k = k;
                params.similarity_threshold = -2.0f; // Include all results
                auto results = db.search(query, params);

                auto end = std::chrono::high_resolution_clock::now();
                auto duration_us =
                    std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
                latencies_us.push_back(static_cast<double>(duration_us));
            }

            // Calculate statistics
            std::sort(latencies_us.begin(), latencies_us.end());

            HNSWBenchmarkResult result;
            result.name = "HNSW_Search_" + std::to_string(corpus_size);
            result.corpus_size = corpus_size;
            result.dimension = dim;
            result.k = k;
            result.mean_latency_us =
                std::accumulate(latencies_us.begin(), latencies_us.end(), 0.0) /
                latencies_us.size();
            result.p50_latency_us = latencies_us[latencies_us.size() / 2];
            result.p95_latency_us = latencies_us[latencies_us.size() * 95 / 100];
            result.p99_latency_us = latencies_us[latencies_us.size() * 99 / 100];
            result.qps = 1000000.0 / result.mean_latency_us;

            // Calculate recall
            VectorSearchParams params;
            params.k = k;
            params.similarity_threshold = -2.0f; // Include all results
            auto final_results = db.search(query, params);
            std::vector<std::string> retrieved;
            for (const auto& r : final_results) {
                retrieved.push_back(r.chunk_id);
            }
            result.recall = calculateRecall(retrieved, gt_indices, chunk_ids);

            all_results.push_back(result);

            INFO("  Mean latency: " << result.mean_latency_us << " μs");
            INFO("  P95 latency: " << result.p95_latency_us << " μs");
            INFO("  QPS: " << result.qps);
            INFO("  Recall@" << k << ": " << result.recall);
            INFO("  Retrieved " << retrieved.size() << " results");

            // Performance requirements (relaxed for TSan builds)
            // Normal: <10ms, TSan: <100ms
            CHECK(result.mean_latency_us < 100000); // <100ms mean latency (TSan overhead)
            CHECK(result.recall >= 0.5);            // >50% recall (approximate search)
        }

        writeBenchmarkJSON("bench_results/hnsw_search_corpus_scaling.json", all_results);
    }

    SECTION("Search Latency vs k") {
        size_t corpus_size = 10000;
        size_t dim = 384;
        std::vector<size_t> k_values = {1, 5, 10, 20, 50, 100};

        auto config = createConfig(dim);
        VectorDatabase db(config);
        REQUIRE(db.initialize());

        auto vectors = generateRandomVectors(corpus_size, dim, rng);
        std::vector<std::string> chunk_ids;

        for (size_t i = 0; i < corpus_size; ++i) {
            std::string chunk_id = "chunk_" + std::to_string(i);
            chunk_ids.push_back(chunk_id);
            db.insertVector(createRecord(chunk_id, "doc_" + std::to_string(i), vectors[i]));
        }

        auto query = generateRandomVector(dim, rng);

        std::vector<HNSWBenchmarkResult> k_results;

        for (size_t k : k_values) {
            INFO("Testing k=" << k);

            auto gt_indices = bruteForceKNN(query, vectors, k);

            std::vector<double> latencies_us;
            size_t num_iterations = 100;

            for (size_t iter = 0; iter < num_iterations; ++iter) {
                auto start = std::chrono::high_resolution_clock::now();

                VectorSearchParams params;
                params.k = k;
                params.similarity_threshold = -2.0f; // Include all results
                auto results = db.search(query, params);

                auto end = std::chrono::high_resolution_clock::now();
                auto duration_us =
                    std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
                latencies_us.push_back(static_cast<double>(duration_us));
            }

            std::sort(latencies_us.begin(), latencies_us.end());

            HNSWBenchmarkResult result;
            result.name = "HNSW_Search_k" + std::to_string(k);
            result.corpus_size = corpus_size;
            result.dimension = dim;
            result.k = k;
            result.mean_latency_us =
                std::accumulate(latencies_us.begin(), latencies_us.end(), 0.0) /
                latencies_us.size();
            result.p50_latency_us = latencies_us[latencies_us.size() / 2];
            result.p95_latency_us = latencies_us[latencies_us.size() * 95 / 100];
            result.p99_latency_us = latencies_us[latencies_us.size() * 99 / 100];
            result.qps = 1000000.0 / result.mean_latency_us;

            VectorSearchParams params;
            params.k = k;
            params.similarity_threshold = -2.0f; // Include all results
            auto final_results = db.search(query, params);
            std::vector<std::string> retrieved;
            for (const auto& r : final_results) {
                retrieved.push_back(r.chunk_id);
            }
            result.recall = calculateRecall(retrieved, gt_indices, chunk_ids);

            k_results.push_back(result);

            INFO("  k=" << k << " Mean latency: " << result.mean_latency_us << " μs");
            INFO("  Recall: " << result.recall << " Retrieved: " << retrieved.size());
        }

        writeBenchmarkJSON("bench_results/hnsw_search_k_scaling.json", k_results);
    }

    SECTION("Search Latency vs Dimension") {
        size_t corpus_size = 10000;
        std::vector<size_t> dimensions = {64, 128, 256, 384, 512, 768};
        size_t k = 10;

        std::vector<HNSWBenchmarkResult> dim_results;

        for (size_t dim : dimensions) {
            INFO("Testing dimension: " << dim);

            auto config = createConfig(dim);
            VectorDatabase db(config);
            REQUIRE(db.initialize());

            auto vectors = generateRandomVectors(corpus_size, dim, rng);
            std::vector<std::string> chunk_ids;

            for (size_t i = 0; i < corpus_size; ++i) {
                std::string chunk_id = "chunk_" + std::to_string(i);
                chunk_ids.push_back(chunk_id);
                db.insertVector(createRecord(chunk_id, "doc_" + std::to_string(i), vectors[i]));
            }

            auto query = generateRandomVector(dim, rng);
            auto gt_indices = bruteForceKNN(query, vectors, k);

            std::vector<double> latencies_us;
            size_t num_iterations = 100;

            for (size_t iter = 0; iter < num_iterations; ++iter) {
                auto start = std::chrono::high_resolution_clock::now();

                VectorSearchParams params;
                params.k = k;
                params.similarity_threshold = -2.0f; // Include all results
                auto results = db.search(query, params);

                auto end = std::chrono::high_resolution_clock::now();
                auto duration_us =
                    std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
                latencies_us.push_back(static_cast<double>(duration_us));
            }

            std::sort(latencies_us.begin(), latencies_us.end());

            HNSWBenchmarkResult result;
            result.name = "HNSW_Search_dim" + std::to_string(dim);
            result.corpus_size = corpus_size;
            result.dimension = dim;
            result.k = k;
            result.mean_latency_us =
                std::accumulate(latencies_us.begin(), latencies_us.end(), 0.0) /
                latencies_us.size();
            result.p50_latency_us = latencies_us[latencies_us.size() / 2];
            result.p95_latency_us = latencies_us[latencies_us.size() * 95 / 100];
            result.p99_latency_us = latencies_us[latencies_us.size() * 99 / 100];
            result.qps = 1000000.0 / result.mean_latency_us;

            VectorSearchParams params;
            params.k = k;
            params.similarity_threshold = -2.0f; // Include all results
            auto final_results = db.search(query, params);
            std::vector<std::string> retrieved;
            for (const auto& r : final_results) {
                retrieved.push_back(r.chunk_id);
            }
            result.recall = calculateRecall(retrieved, gt_indices, chunk_ids);

            dim_results.push_back(result);

            INFO("  dim=" << dim << " Mean latency: " << result.mean_latency_us << " μs");
            INFO("  Recall: " << result.recall << " Retrieved: " << retrieved.size());
        }

        writeBenchmarkJSON("bench_results/hnsw_search_dim_scaling.json", dim_results);
    }
}

// ============================================================================
// Entity Vector Search Benchmark (Symbol Extractor Integration)
// ============================================================================

TEST_CASE_METHOD(HNSWSearchFixture, "Entity Vector Search with Symbol Extractor",
                 "[benchmark][hnsw][entity][symbol_extractor][!benchmark]") {
    auto plugin = loadSymbolExtractorPlugin();
    if (!plugin.has_value()) {
        WARN("Symbol extractor plugin not found, skipping entity search benchmark");
        return;
    }

    std::mt19937 rng(42);
    size_t dim = 384; // Common embedding dimension

    SECTION("Extract Symbols and Search") {
        // Collect source files (reduced for faster testing)
        auto files = collectSourceFiles("src", 20);
        if (files.empty()) {
            files = collectSourceFiles(".", 20);
        }
        REQUIRE(!files.empty());

        INFO("Found " << files.size() << " source files");

        // Extract symbols from files
        auto extract_start = std::chrono::high_resolution_clock::now();

        std::vector<std::pair<std::string, std::string>> all_symbols;
        for (const auto& file : files) {
            auto symbols = extractSymbols(*plugin, file);
            all_symbols.insert(all_symbols.end(), symbols.begin(), symbols.end());
        }

        auto extract_end = std::chrono::high_resolution_clock::now();
        double extraction_time_ms = static_cast<double>(
            std::chrono::duration_cast<std::chrono::milliseconds>(extract_end - extract_start)
                .count());

        INFO("Extracted " << all_symbols.size() << " symbols in " << extraction_time_ms << " ms");
        REQUIRE(!all_symbols.empty());

        // Create vector database and index symbols with random embeddings
        // (In production, these would be generated by an embedding model)
        auto config = createConfig(dim);
        VectorDatabase db(config);
        REQUIRE(db.initialize());

        auto index_start = std::chrono::high_resolution_clock::now();

        std::vector<std::string> entity_ids;
        std::vector<std::vector<float>> entity_vectors;

        for (size_t i = 0; i < all_symbols.size(); ++i) {
            const auto& [name, qname] = all_symbols[i];
            std::string entity_id = "entity_" + std::to_string(i) + "_" + name;
            entity_ids.push_back(entity_id);

            auto vec = generateRandomVector(dim, rng);
            entity_vectors.push_back(vec);

            VectorRecord rec;
            rec.chunk_id = entity_id;
            rec.document_hash = "entity_doc";
            rec.embedding = vec;
            rec.content = qname;
            rec.start_offset = 0;
            rec.end_offset = qname.size();
            db.insertVector(rec);
        }

        auto index_end = std::chrono::high_resolution_clock::now();
        double indexing_time_ms = static_cast<double>(
            std::chrono::duration_cast<std::chrono::milliseconds>(index_end - index_start).count());

        INFO("Indexed " << all_symbols.size() << " entity vectors in " << indexing_time_ms
                        << " ms");

        // Benchmark search
        auto query = generateRandomVector(dim, rng);
        std::vector<double> latencies_us;
        size_t num_iterations = 50; // Reduced for faster testing
        size_t k = 10;

        for (size_t iter = 0; iter < num_iterations; ++iter) {
            auto start = std::chrono::high_resolution_clock::now();

            VectorSearchParams params;
            params.k = k;
            params.similarity_threshold = -2.0f; // Include all results
            auto results = db.search(query, params);

            auto end = std::chrono::high_resolution_clock::now();
            auto duration_us =
                std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
            latencies_us.push_back(static_cast<double>(duration_us));
        }

        std::sort(latencies_us.begin(), latencies_us.end());
        double mean_latency = latencies_us.empty()
                                  ? 0.0
                                  : std::accumulate(latencies_us.begin(), latencies_us.end(), 0.0) /
                                        latencies_us.size();

        EntitySearchResult result;
        result.name = "Entity_Vector_Search";
        result.total_entities = all_symbols.size();
        result.total_files = files.size();
        result.extraction_time_ms = extraction_time_ms;
        result.indexing_time_ms = indexing_time_ms;
        result.search_latency_us = mean_latency;
        result.qps = mean_latency > 0.0 ? 1000000.0 / mean_latency : 0.0;

        INFO("Entity Search Results:");
        INFO("  Total entities: " << result.total_entities);
        INFO("  Total files: " << result.total_files);
        INFO("  Extraction time: " << result.extraction_time_ms << " ms");
        INFO("  Indexing time: " << result.indexing_time_ms << " ms");
        INFO("  Search latency: " << result.search_latency_us << " μs");
        INFO("  QPS: " << result.qps);

        writeEntitySearchJSON("bench_results/entity_vector_search.json", result);

        // Performance requirements (relaxed for TSan builds)
        CHECK(result.search_latency_us < 50000); // <50ms search latency (TSan overhead)
        CHECK(result.qps > 20);                  // >20 QPS (TSan overhead)
    }
}

// ============================================================================
// Batch Insert Performance
// ============================================================================

TEST_CASE_METHOD(HNSWSearchFixture, "HNSW Batch Insert Performance",
                 "[benchmark][hnsw][insert][!benchmark]") {
    std::mt19937 rng(42);
    size_t dim = 384;

    SECTION("Batch Insert Throughput") {
        std::vector<size_t> batch_sizes = {10, 50, 100, 500, 1000};
        size_t total_vectors = 10000;

        json results = json::array();

        for (size_t batch_size : batch_sizes) {
            INFO("Testing batch size: " << batch_size);

            auto config = createConfig(dim);
            VectorDatabase db(config);
            REQUIRE(db.initialize());

            auto vectors = generateRandomVectors(total_vectors, dim, rng);

            auto start = std::chrono::high_resolution_clock::now();

            for (size_t i = 0; i < total_vectors; i += batch_size) {
                std::vector<VectorRecord> batch;
                for (size_t j = i; j < std::min(i + batch_size, total_vectors); ++j) {
                    batch.push_back(createRecord("chunk_" + std::to_string(j),
                                                 "doc_" + std::to_string(j), vectors[j]));
                }
                db.insertVectorsBatch(batch);
            }

            auto end = std::chrono::high_resolution_clock::now();
            double duration_ms = static_cast<double>(
                std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count());
            double throughput = total_vectors / (duration_ms / 1000.0);

            INFO("  Batch size " << batch_size << ": " << duration_ms << " ms, " << throughput
                                 << " vectors/sec");

            results.push_back({{"batch_size", batch_size},
                               {"total_vectors", total_vectors},
                               {"duration_ms", duration_ms},
                               {"throughput_vectors_per_sec", throughput}});

            CHECK(throughput > 10); // >10 vectors/sec minimum (relaxed for TSan)
        }

        fs::create_directories("bench_results");
        std::ofstream out("bench_results/hnsw_batch_insert.json");
        out << results.dump(2);
    }
}