#include "qkmeans/anns/faiss_backend.hpp"
#include <stdexcept>
#include <limits>

#ifdef QKMEANS_HAS_FAISS

namespace qkmeans {
namespace anns {

FAISSIndex::FAISSIndex(const ANNSConfig& config)
    : dim_(config.dimension)
    , index_type_(static_cast<FAISSIndexType>(config.faiss_index_type))
    , hnsw_M_(config.faiss_hnsw_M)
    , hnsw_efConstruction_(config.faiss_hnsw_efConstruction)
    , hnsw_efSearch_(config.faiss_hnsw_efSearch)
    , lsh_nbits_(config.faiss_lsh_nbits)
    , ivf_nlist_(config.faiss_nlist)
    , ivf_nprobe_(config.faiss_nprobe)
    , query_buffer_(config.dimension)
    , insert_buffer_(config.dimension)
{
    if (dim_ == 0) {
        throw std::invalid_argument("Dimension must be positive");
    }
    create_index();
}

FAISSIndex::FAISSIndex(size_t dim, FAISSIndexType index_type,
                       size_t hnsw_M, size_t hnsw_efConstruction, size_t hnsw_efSearch,
                       size_t lsh_nbits, size_t ivf_nlist, size_t ivf_nprobe)
    : dim_(dim)
    , index_type_(index_type)
    , hnsw_M_(hnsw_M)
    , hnsw_efConstruction_(hnsw_efConstruction)
    , hnsw_efSearch_(hnsw_efSearch)
    , lsh_nbits_(lsh_nbits == 0 ? dim * 2 : lsh_nbits)
    , ivf_nlist_(ivf_nlist)
    , ivf_nprobe_(ivf_nprobe)
    , query_buffer_(dim)
    , insert_buffer_(dim)
{
    if (dim_ == 0) {
        throw std::invalid_argument("Dimension must be positive");
    }
    create_index();
}

FAISSIndex::~FAISSIndex() = default;

void FAISSIndex::create_index() {
    switch (index_type_) {
        case FAISSIndexType::FLAT:
            index_ = std::make_unique<faiss::IndexFlatL2>(dim_);
            break;

        case FAISSIndexType::HNSW: {
            auto hnsw = std::make_unique<faiss::IndexHNSWFlat>(dim_, hnsw_M_);
            hnsw->hnsw.efConstruction = hnsw_efConstruction_;
            hnsw->hnsw.efSearch = hnsw_efSearch_;
            index_ = std::move(hnsw);
            break;
        }

        case FAISSIndexType::LSH:
            index_ = std::make_unique<faiss::IndexLSH>(dim_, lsh_nbits_);
            break;

        case FAISSIndexType::IVF:
            // For IVF, we use the flat quantizer directly until we have enough data to train
            ivf_quantizer_ = std::make_unique<faiss::IndexFlatL2>(dim_);
            index_ = std::make_unique<faiss::IndexFlatL2>(dim_);  // Start with flat
            ivf_trained_ = false;
            break;

        default:
            throw std::invalid_argument("Unknown FAISS index type");
    }
}

void FAISSIndex::insert(const double* point, size_t /* idx */) {
    // Convert double to float
    for (size_t i = 0; i < dim_; ++i) {
        insert_buffer_[i] = static_cast<float>(point[i]);
    }

    // For IVF, we need to collect training data
    if (index_type_ == FAISSIndexType::IVF && !ivf_trained_) {
        training_data_.insert(training_data_.end(), insert_buffer_.begin(), insert_buffer_.end());

        // Train when we have enough data
        if (training_data_.size() / dim_ >= ivf_nlist_ * 10) {
            // Create and train the actual IVF index
            auto ivf_index = std::make_unique<faiss::IndexIVFFlat>(
                ivf_quantizer_.get(), dim_, ivf_nlist_);
            ivf_index->train(training_data_.size() / dim_, training_data_.data());
            ivf_index->add(training_data_.size() / dim_, training_data_.data());
            ivf_index->nprobe = ivf_nprobe_;
            index_ = std::move(ivf_index);
            ivf_trained_ = true;
            num_vectors_ = training_data_.size() / dim_;
            training_data_.clear();
            return;
        }
    }

    // Add to index
    index_->add(1, insert_buffer_.data());
    ++num_vectors_;
}

QueryResult FAISSIndex::query(const double* point) const {
    if (num_vectors_ == 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]);
    }

    faiss::idx_t label;
    float distance;

    // Search for nearest neighbor
    index_->search(1, query_buffer_.data(), 1, &distance, &label);

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

    return {static_cast<size_t>(label), static_cast<double>(distance)};
}

void FAISSIndex::batch_query(const double* points, size_t n,
                             size_t* indices_out, double* distances_out) const {
    if (num_vectors_ == 0) {
        for (size_t i = 0; i < n; ++i) {
            indices_out[i] = 0;
            distances_out[i] = std::numeric_limits<double>::max();
        }
        return;
    }

    // Convert all points to float
    std::vector<float> points_float(n * dim_);
    for (size_t i = 0; i < n * dim_; ++i) {
        points_float[i] = static_cast<float>(points[i]);
    }

    std::vector<faiss::idx_t> labels(n);
    std::vector<float> distances(n);

    // Batch search
    index_->search(n, points_float.data(), 1, distances.data(), labels.data());

    // Copy results
    for (size_t i = 0; i < n; ++i) {
        indices_out[i] = labels[i] >= 0 ? static_cast<size_t>(labels[i]) : 0;
        distances_out[i] = static_cast<double>(distances[i]);
    }
}

size_t FAISSIndex::size() const {
    return num_vectors_;
}

void FAISSIndex::clear() {
    create_index();
    num_vectors_ = 0;
    training_data_.clear();
    ivf_trained_ = false;
}

void FAISSIndex::reserve(size_t /* n */) {
    // Most FAISS indexes don't support pre-reservation
}

} // namespace anns
} // namespace qkmeans

#endif // QKMEANS_HAS_FAISS
