#include "qkmeans/anns/hnsw_backend.hpp"
#include <stdexcept>

#ifdef QKMEANS_HAS_HNSWLIB

namespace qkmeans {
namespace anns {

HNSWIndex::HNSWIndex(const ANNSConfig& config)
    : HNSWIndex(config.dimension, config.expected_size,
                config.hnsw_M, config.hnsw_ef_construction, config.hnsw_ef_search)
{
}

HNSWIndex::HNSWIndex(size_t dim, size_t max_elements,
                     size_t M, size_t ef_construction, size_t ef_search)
    : dim_(dim)
    , max_elements_(max_elements)
    , M_(M)
    , ef_construction_(ef_construction)
    , ef_search_(ef_search)
    , query_buffer_(dim)
    , insert_buffer_(dim)
{
    if (dim == 0) {
        throw std::invalid_argument("Dimension must be positive");
    }
    if (max_elements == 0) {
        max_elements_ = 10000;  // Default max elements
    }

    // Create L2 space (hnswlib uses float)
    space_ = std::make_unique<hnswlib::L2Space>(dim);

    // Create HNSW index
    index_ = std::make_unique<hnswlib::HierarchicalNSW<float>>(
        space_.get(),
        max_elements_,
        M_,
        ef_construction_
    );

    // Set search parameter
    index_->setEf(ef_search_);
}

void HNSWIndex::insert(const double* point, size_t idx) {
    // Resize index if needed
    if (num_points_ >= max_elements_) {
        max_elements_ *= 2;
        index_->resizeIndex(max_elements_);
    }

    // Convert double to float
    for (size_t i = 0; i < dim_; ++i) {
        insert_buffer_[i] = static_cast<float>(point[i]);
    }

    // Add to index with sequential label (we track separately)
    index_->addPoint(insert_buffer_.data(), num_points_);
    ++num_points_;
}

QueryResult HNSWIndex::query(const double* point) const {
    if (num_points_ == 0) {
        return {0, std::numeric_limits<double>::max()};
    }

    // Convert double to float
    for (size_t i = 0; i < dim_; ++i) {
        query_buffer_[i] = static_cast<float>(point[i]);
    }

    // Search for nearest neighbor
    auto result = index_->searchKnn(query_buffer_.data(), 1);

    if (result.empty()) {
        return {0, std::numeric_limits<double>::max()};
    }

    // Result is a priority queue with (distance, label) pairs
    auto& top = result.top();
    size_t label = top.second;
    double distance_sq = static_cast<double>(top.first);

    return {label, distance_sq};
}

void HNSWIndex::batch_query(const double* points, size_t n,
                            size_t* indices_out, double* distances_out) const {
    // Use OpenMP for parallel queries
    #pragma omp parallel for
    for (size_t i = 0; i < n; ++i) {
        // Each thread needs its own buffer
        std::vector<float> local_buffer(dim_);
        for (size_t j = 0; j < dim_; ++j) {
            local_buffer[j] = static_cast<float>(points[i * dim_ + j]);
        }

        if (num_points_ == 0) {
            indices_out[i] = 0;
            distances_out[i] = std::numeric_limits<double>::max();
        } else {
            auto result = index_->searchKnn(local_buffer.data(), 1);
            if (!result.empty()) {
                auto& top = result.top();
                indices_out[i] = top.second;
                distances_out[i] = static_cast<double>(top.first);
            } else {
                indices_out[i] = 0;
                distances_out[i] = std::numeric_limits<double>::max();
            }
        }
    }
}

size_t HNSWIndex::size() const {
    return num_points_;
}

void HNSWIndex::clear() {
    // hnswlib doesn't have a clear method, so we recreate the index
    index_ = std::make_unique<hnswlib::HierarchicalNSW<float>>(
        space_.get(),
        max_elements_,
        M_,
        ef_construction_
    );
    index_->setEf(ef_search_);
    num_points_ = 0;
}

void HNSWIndex::reserve(size_t n) {
    if (n > max_elements_) {
        max_elements_ = n;
        // Note: resizeIndex only works after initialization
        if (index_ && num_points_ == 0) {
            index_ = std::make_unique<hnswlib::HierarchicalNSW<float>>(
                space_.get(),
                max_elements_,
                M_,
                ef_construction_
            );
            index_->setEf(ef_search_);
        }
    }
}

} // namespace anns
} // namespace qkmeans

#endif // QKMEANS_HAS_HNSWLIB
