#include "common_includes.h"
#include "parameters.h"

namespace diskann
{
enum DataStoreStrategy
{
    MEMORY
};

enum GraphStoreStrategy
{
};
struct IndexConfig
{
    DataStoreStrategy data_strategy;
    GraphStoreStrategy graph_strategy;

    Metric metric;
    size_t dimension;
    size_t max_points;

    bool dynamic_index;
    bool enable_tags;
    bool pq_dist_build;
    bool concurrent_consolidate;
    bool use_opq;

    size_t num_pq_chunks;
    size_t num_frozen_pts;

    std::string label_type;
    std::string tag_type;
    std::string data_type;

    std::shared_ptr<IndexWriteParameters> index_write_params;

    uint32_t search_threads;
    uint32_t initial_search_list_size;

  private:
    IndexConfig(DataStoreStrategy data_strategy, GraphStoreStrategy graph_strategy, Metric metric, size_t dimension,
                size_t max_points, size_t num_pq_chunks, size_t num_frozen_points, bool dynamic_index, bool enable_tags,
                bool pq_dist_build, bool concurrent_consolidate, bool use_opq, const std::string &data_type,
                const std::string &tag_type, const std::string &label_type,
                std::shared_ptr<IndexWriteParameters> index_write_params, uint32_t search_threads,
                uint32_t initial_search_list_size)
        : data_strategy(data_strategy), graph_strategy(graph_strategy), metric(metric), dimension(dimension),
          max_points(max_points), dynamic_index(dynamic_index), enable_tags(enable_tags), pq_dist_build(pq_dist_build),
          concurrent_consolidate(concurrent_consolidate), use_opq(use_opq), num_pq_chunks(num_pq_chunks),
          num_frozen_pts(num_frozen_points), label_type(label_type), tag_type(tag_type), data_type(data_type),
          index_write_params(index_write_params), search_threads(search_threads),
          initial_search_list_size(initial_search_list_size)
    {
    }

    friend class IndexConfigBuilder;
};

class IndexConfigBuilder
{
  public:
    IndexConfigBuilder()
    {
    }

    IndexConfigBuilder &with_metric(Metric m)
    {
        this->_metric = m;
        return *this;
    }

    IndexConfigBuilder &with_graph_load_store_strategy(GraphStoreStrategy graph_strategy)
    {
        this->_graph_strategy = graph_strategy;
        return *this;
    }

    IndexConfigBuilder &with_data_load_store_strategy(DataStoreStrategy data_strategy)
    {
        this->_data_strategy = data_strategy;
        return *this;
    }

    IndexConfigBuilder &with_dimension(size_t dimension)
    {
        this->_dimension = dimension;
        return *this;
    }

    IndexConfigBuilder &with_max_points(size_t max_points)
    {
        this->_max_points = max_points;
        return *this;
    }

    IndexConfigBuilder &is_dynamic_index(bool dynamic_index)
    {
        this->_dynamic_index = dynamic_index;
        return *this;
    }

    IndexConfigBuilder &is_enable_tags(bool enable_tags)
    {
        this->_enable_tags = enable_tags;
        return *this;
    }

    IndexConfigBuilder &is_pq_dist_build(bool pq_dist_build)
    {
        this->_pq_dist_build = pq_dist_build;
        return *this;
    }

    IndexConfigBuilder &is_concurrent_consolidate(bool concurrent_consolidate)
    {
        this->_concurrent_consolidate = concurrent_consolidate;
        return *this;
    }

    IndexConfigBuilder &is_use_opq(bool use_opq)
    {
        this->_use_opq = use_opq;
        return *this;
    }

    IndexConfigBuilder &with_num_pq_chunks(size_t num_pq_chunks)
    {
        this->_num_pq_chunks = num_pq_chunks;
        return *this;
    }

    IndexConfigBuilder &with_num_frozen_pts(size_t num_frozen_pts)
    {
        this->_num_frozen_pts = num_frozen_pts;
        return *this;
    }

    IndexConfigBuilder &with_label_type(const std::string &label_type)
    {
        this->_label_type = label_type;
        return *this;
    }

    IndexConfigBuilder &with_tag_type(const std::string &tag_type)
    {
        this->_tag_type = tag_type;
        return *this;
    }

    IndexConfigBuilder &with_data_type(const std::string &data_type)
    {
        this->_data_type = data_type;
        return *this;
    }

    IndexConfigBuilder &with_index_write_params(IndexWriteParameters &index_write_params)
    {
        this->_index_write_params = std::make_shared<IndexWriteParameters>(index_write_params);
        return *this;
    }

    IndexConfigBuilder &with_search_threads(uint32_t search_threads)
    {
        this->_search_threads = search_threads;
        return *this;
    }

    IndexConfigBuilder &with_initial_search_list_size(uint32_t search_list_size)
    {
        this->_initial_search_list_size = search_list_size;
        return *this;
    }

    IndexConfig build()
    {
        if (_data_type == "" || _data_type.empty())
            throw ANNException("Error: data_type can not be empty", -1);

        if (_dynamic_index && _index_write_params != nullptr)
        {
            if (_search_threads == 0)
                throw ANNException("Error: please pass search_threads for building dynamic index.", -1);

            if (_initial_search_list_size == 0)
                throw ANNException("Error: please pass initial_search_list_size for building dynamic index.", -1);
        }

        return IndexConfig(_data_strategy, _graph_strategy, _metric, _dimension, _max_points, _num_pq_chunks,
                           _num_frozen_pts, _dynamic_index, _enable_tags, _pq_dist_build, _concurrent_consolidate,
                           _use_opq, _data_type, _tag_type, _label_type, _index_write_params, _search_threads,
                           _initial_search_list_size);
    }

    IndexConfigBuilder(const IndexConfigBuilder &) = delete;
    IndexConfigBuilder &operator=(const IndexConfigBuilder &) = delete;

  private:
    DataStoreStrategy _data_strategy;
    GraphStoreStrategy _graph_strategy;

    Metric _metric;
    size_t _dimension;
    size_t _max_points;

    bool _dynamic_index = false;
    bool _enable_tags = false;
    bool _pq_dist_build = false;
    bool _concurrent_consolidate = false;
    bool _use_opq = false;

    size_t _num_pq_chunks = 0;
    size_t _num_frozen_pts = 0;

    std::string _label_type = "uint32";
    std::string _tag_type = "uint32";
    std::string _data_type;

    std::shared_ptr<IndexWriteParameters> _index_write_params;

    uint32_t _search_threads;
    uint32_t _initial_search_list_size;
};
} // namespace diskann
