#pragma once

#include "base.hpp"

#ifdef QKMEANS_HAS_FAISS
#include <faiss/IndexFlat.h>
#include <faiss/IndexHNSW.h>
#include <faiss/IndexLSH.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/Index.h>
#endif

namespace qkmeans {
namespace anns {

// FAISS index types
enum class FAISSIndexType {
    FLAT,       // IndexFlatL2 - exact brute force
    HNSW,       // IndexHNSWFlat - HNSW graph
    LSH,        // IndexLSH - locality-sensitive hashing
    IVF         // IndexIVFFlat - inverted file index
};

#ifdef QKMEANS_HAS_FAISS

/**
 * @brief FAISS-based approximate nearest neighbor index.
 *
 * Supports multiple FAISS index types:
 * - FLAT: IndexFlatL2 (exact brute force)
 * - HNSW: IndexHNSWFlat (graph-based approximate)
 * - LSH: IndexLSH (locality-sensitive hashing)
 * - IVF: IndexIVFFlat (inverted file with flat quantizer)
 */
class FAISSIndex : public ANNSBase {
public:
    /**
     * @brief Construct FAISS index from config.
     */
    explicit FAISSIndex(const ANNSConfig& config);

    /**
     * @brief Construct FAISS index with explicit parameters.
     */
    FAISSIndex(size_t dim, FAISSIndexType index_type = FAISSIndexType::FLAT,
               size_t hnsw_M = 32, size_t hnsw_efConstruction = 40, size_t hnsw_efSearch = 16,
               size_t lsh_nbits = 0,  // 0 = auto (2*dim)
               size_t ivf_nlist = 100, size_t ivf_nprobe = 10);

    ~FAISSIndex() override;

    void insert(const double* point, size_t idx) override;
    QueryResult query(const double* point) const override;
    void batch_query(const double* points, size_t n,
                     size_t* indices_out, double* distances_out) const override;
    size_t size() const override;
    size_t dimension() const override { return dim_; }
    void clear() override;
    void reserve(size_t n) override;

private:
    void create_index();

    size_t dim_;
    FAISSIndexType index_type_;

    // HNSW parameters
    size_t hnsw_M_;
    size_t hnsw_efConstruction_;
    size_t hnsw_efSearch_;

    // LSH parameters
    size_t lsh_nbits_;

    // IVF parameters
    size_t ivf_nlist_;
    size_t ivf_nprobe_;

    // The actual FAISS index (type erased)
    std::unique_ptr<faiss::Index> index_;

    // For IVF, we need to keep the quantizer alive
    std::unique_ptr<faiss::IndexFlatL2> ivf_quantizer_;

    // Buffer for float conversion
    mutable std::vector<float> query_buffer_;
    std::vector<float> insert_buffer_;
    size_t num_vectors_ = 0;

    // For IVF training
    std::vector<float> training_data_;
    bool ivf_trained_ = false;
};

#endif // QKMEANS_HAS_FAISS

} // namespace anns
} // namespace qkmeans
