// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "bench_result.h"
#include "helper.h"

static bool g_debug_mode = 0;

//------------------------------------------------------------
// Bench
//------------------------------------------------------------
enum RetrievalMode { RM_UNDEFINED = 0, RM_DENSE = 1, RM_SPARSE = 2 };

enum FilterMode { FM_UNDEFINED = 0, FM_NONE = 1, FM_TAG = 2 };

template <typename T>
class Bench {
 public:
  Bench(size_t threads, size_t bench_secs, size_t batch_count,
        RetrievalMode &retrieval_mode, FilterMode filter_mode)
      : threads_(threads),
        bench_secs_(bench_secs),
        batch_count_(batch_count),
        retrieval_mode_{retrieval_mode},
        filter_mode_{filter_mode} {
    if (threads_ == 0) {
      pool_ = make_shared<ThreadPool>(false);
      threads_ = pool_->count();
      cout << "Using cpu count as thread pool count[" << threads_ << "]"
           << endl;
    } else {
      pool_ = make_shared<ThreadPool>(threads_, false);
      cout << "Using thread pool count[" << threads_ << "]" << endl;
    }
    if (batch_count_ < 1) {
      batch_count_ = 1;
    }
  }

  static void stop(int signo) {
    if (STOP_NOW) {
      exit(signo);
    }
    STOP_NOW = true;
    cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
         << flush;
  }

  bool load_query(const std::string &query_file, const std::string &first_sep,
                  const std::string &second_sep) {
    TxtInputReader<T> reader;
    vector<vector<T>> queries;
    vector<SparseData<T>> sparse_data;
    vector<vector<uint64_t>> taglists;

    if (!reader.load_query(query_file, first_sep, second_sep, queries,
                           sparse_data, taglists)) {
      LOG_ERROR("Load query error");
      return false;
    }

    if (batch_count_ == 1) {
      batch_queries_ = queries;

      for (size_t i = 0; i < sparse_data.size(); ++i) {
        vector<uint32_t> sparse_count;
        sparse_count.push_back(sparse_data[i].count);

        batch_sparse_counts_.push_back(sparse_count);
        batch_sparse_indices_.push_back(sparse_data[i].indices);
        batch_sparse_features_.push_back(sparse_data[i].features);
      }

      for (size_t i = 0; i < taglists.size(); ++i) {
        vector<vector<uint64_t>> new_taglists;
        new_taglists.push_back(taglists[i]);

        batch_taglists_.push_back(std::move(new_taglists));
      }
    } else {
      size_t num_batch = (queries.size() + batch_count_ - 1) / batch_count_;
      size_t idx = 0;
      for (size_t n = 0; n < num_batch; ++n) {
        vector<T> batch_query;
        vector<uint32_t> batch_sparse_count;
        vector<uint32_t> batch_sparse_indices;
        vector<T> batch_sparse_feature;
        vector<vector<uint64_t>> batch_taglists;

        for (size_t i = 0; i < batch_count_; ++i) {
          for (size_t k = 0; k < queries[idx].size(); ++k) {
            batch_query.push_back(queries[idx][k]);
          }

          batch_sparse_count.push_back(sparse_data[idx].count);

          for (size_t k = 0; k < sparse_data[idx].indices.size(); ++k) {
            batch_sparse_indices.push_back(sparse_data[idx].indices[k]);
          }

          for (size_t k = 0; k < sparse_data[idx].features.size(); ++k) {
            batch_sparse_feature.push_back(sparse_data[idx].features[k]);
          }

          if (taglists.size() > idx) {
            batch_taglists.push_back(taglists[idx]);
          }

          idx = (idx + 1) % queries.size();
        }

        batch_queries_.push_back(batch_query);
        batch_sparse_counts_.push_back(batch_sparse_count);
        batch_sparse_indices_.push_back(batch_sparse_indices);
        batch_sparse_features_.push_back(batch_sparse_feature);
        batch_taglists_.push_back(batch_taglists);
      }
    }

    dim_ = queries[0].size();
    if (typeid(T) == typeid(float)) {
      qmeta_.set_meta(IndexMeta::DataType::DT_FP32, dim_);
    } else if (typeid(T) == typeid(int8_t)) {
      qmeta_.set_meta(IndexMeta::DataType::DT_INT8, dim_);
    } else {
      LOG_ERROR("unsupported type");
      return false;
    }

    cout << "Load query done!" << endl;
    return true;
  }

  void run(core_interface::Index::Pointer index,
           core_interface::BaseIndexQueryParam::Pointer query_param,
           int max_iter, int topk) {
    // Check
    if (batch_queries_.size() == 0) {
      return;
    }

    query_param_ = query_param;
    query_param_->topk = topk;
    query_param_->is_linear = false;

    // Do bench
    signal(SIGINT, stop);
    bench_result_.mark_start();
    auto start_time = Monotime::MilliSeconds();
    for (size_t i = 0; i < threads_; ++i) {
      pool_->execute(this, &Bench<T>::start_bench, index, max_iter, &STOP_NOW);
    }

    while (!pool_->is_finished()) {
      this_thread::sleep_for(chrono::milliseconds(1));
      if (Monotime::MilliSeconds() - start_time > bench_secs_ * 1000) {
        STOP_NOW = true;
      }
    }

    pool_->wait_finish();

    bench_result_.mark_end();
    bench_result_.print();
  }

  void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
                     const std::vector<uint64_t> &tag_key_list) {
    id_to_tags_list_ = id_to_tags_list;
    tag_key_list_ = tag_key_list;
  }

 private:
  void start_bench(core_interface::Index::Pointer index, size_t max_iter,
                   const bool *is_stop) {
    size_t thread_index = pool_->indexof_this();

    size_t i = thread_index;
    for (; i < max_iter && !*is_stop; i += threads_) {
      int idx = i % batch_queries_.size();

      // prefilter
      FilterResultCache filter_cache;
      std::shared_ptr<IndexFilter> filter_ptr = nullptr;
      if (filter_mode_ == FM_TAG) {
        if (batch_taglists_[idx].size() != 1) {
          LOG_ERROR("query tag list not equal to one!");
          return;
        }

        int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
                                      tag_key_list_);
        if (ret != 0) {
          LOG_ERROR("prefilter failed, idx: %d", idx);
          return;
        }

        auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };

        filter_ptr = std::make_shared<IndexFilter>();
        filter_ptr->set(filterFunc);
      }

      auto query_param = query_param_->Clone();
      query_param->filter = filter_ptr;


      // Do knn_search
      uint64_t start = Monotime::MicroSeconds();
      int ret;
      if (retrieval_mode_ == RM_DENSE) {
        if (batch_count_ == 1) {
          ret = do_knn_search<T>(index, batch_queries_[idx], query_param);
        } else {
          ret = do_knn_search_batch<T>(index, batch_queries_[idx], query_param);
        }

        if (ret != 0) {
          LOG_ERROR("Failed to knn search, ret=%d %s", ret,
                    IndexError::What(ret));
          return;
        }
      } else {
        std::string mode = retrieval_mode_ == 1 ? "Dense" : "Sparse";
        LOG_ERROR("unsupported retrieval mode: %s", mode.c_str());
      }

      uint64_t end = Monotime::MicroSeconds();

      // Do sample
      bench_result_.add_time(batch_count_, end - start);
    }
  }

  template <typename U>
  typename std::enable_if<
      std::is_same<float, U>::value || std::is_same<int8_t, U>::value ||
          std::is_same<uint32_t, U>::value || std::is_same<uint64_t, U>::value,
      int>::type
  do_knn_search(core_interface::Index::Pointer index, const vector<U> &query,
                core_interface::BaseIndexQueryParam::Pointer query_param) {
    core_interface::DenseVector dense_query;
    dense_query.data = query.data();
    core_interface::VectorData query_data;
    query_data.vector = dense_query;

    core_interface::SearchResult search_result;
    int ret = index->Search(query_data, query_param, &search_result);
    if (ret < 0) {
      return ret;
    }

    if (search_result.doc_list_.empty()) {
      LOG_ERROR("Search results is empty");
    }

    return 0;
  }

  template <typename U>
  typename std::enable_if<
      std::is_same<float, U>::value || std::is_same<int8_t, U>::value ||
          std::is_same<uint32_t, U>::value || std::is_same<uint64_t, U>::value,
      int>::type
  do_knn_search_batch(
      core_interface::Index::Pointer index, const vector<U> &query,
      core_interface::BaseIndexQueryParam::Pointer query_param) {
    // For batch search, we search each query separately
    size_t qnum = query.size() / dim_;
    for (size_t i = 0; i < qnum; ++i) {
      core_interface::DenseVector dense_query;
      dense_query.data = query.data() + i * dim_;
      core_interface::VectorData query_data;
      query_data.vector = dense_query;

      core_interface::SearchResult search_result;
      int ret = index->Search(query_data, query_param, &search_result);
      if (ret < 0) {
        return ret;
      }

      if (search_result.doc_list_.empty()) {
        LOG_ERROR("Search results is empty for batch query %zu", i);
      }
    }

    return 0;
  }

 private:
  IndexQueryMeta qmeta_{};
  size_t threads_;
  size_t bench_secs_;
  size_t batch_count_;
  size_t dim_;
  shared_ptr<ThreadPool> pool_;
  core_interface::BaseIndexQueryParam::Pointer query_param_;

  vector<vector<T>> batch_queries_;
  vector<vector<uint32_t>> batch_sparse_counts_;
  vector<vector<uint32_t>> batch_sparse_indices_;
  vector<vector<T>> batch_sparse_features_;
  vector<vector<vector<uint64_t>>> batch_taglists_;

  // Tag lists for filtering
  std::vector<std::vector<uint64_t>> id_to_tags_list_;
  std::vector<uint64_t> tag_key_list_;

  BenchResult bench_result_;
  RetrievalMode retrieval_mode_{RM_UNDEFINED};
  FilterMode filter_mode_{FM_NONE};
  static bool STOP_NOW;
};

template <typename T>
bool Bench<T>::STOP_NOW = false;

//------------------------------------------------------------
// Sparse Bench
//------------------------------------------------------------
template <typename T>
class SparseBench {
 public:
  SparseBench(size_t threads, size_t bench_secs, size_t batch_count,
              FilterMode filter_mode)
      : threads_(threads),
        bench_secs_(bench_secs),
        batch_count_(batch_count),
        filter_mode_{filter_mode} {
    if (threads_ == 0) {
      pool_ = make_shared<ThreadPool>(false);
      threads_ = pool_->count();
      cout << "Using cpu count as thread pool count[" << threads_ << "]"
           << endl;
    } else {
      pool_ = make_shared<ThreadPool>(threads_, false);
      cout << "Using thread pool count[" << threads_ << "]" << endl;
    }
    if (batch_count_ < 1) {
      batch_count_ = 1;
    }
  }

  static void stop(int signo) {
    if (STOP_NOW) {
      exit(signo);
    }
    STOP_NOW = true;
    cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
         << flush;
  }

  bool load_query(const std::string &query_file, const std::string &first_sep,
                  const std::string &second_sep) {
    TxtInputReader<T> reader;
    vector<vector<T>> queries;
    vector<SparseData<T>> sparse_data;
    vector<vector<uint64_t>> taglists;

    if (!reader.load_query(query_file, first_sep, second_sep, queries,
                           sparse_data, taglists)) {
      LOG_ERROR("Load query error");
      return false;
    }

    linear_sparse_data_ = sparse_data;

    if (batch_count_ == 1) {
      for (size_t i = 0; i < sparse_data.size(); ++i) {
        vector<uint32_t> sparse_count;
        sparse_count.push_back(sparse_data[i].count);

        batch_sparse_counts_.push_back(sparse_count);
        batch_sparse_indices_.push_back(sparse_data[i].indices);
        batch_sparse_features_.push_back(sparse_data[i].features);
      }

      for (size_t i = 0; i < taglists.size(); ++i) {
        vector<vector<uint64_t>> new_taglists;
        new_taglists.push_back(taglists[i]);

        batch_taglists_.push_back(std::move(new_taglists));
      }
    } else {
      size_t num_batch = (queries.size() + batch_count_ - 1) / batch_count_;
      size_t idx = 0;
      for (size_t n = 0; n < num_batch; ++n) {
        vector<uint32_t> batch_sparse_count;
        vector<uint32_t> batch_sparse_indices;
        vector<T> batch_sparse_feature;
        vector<vector<uint64_t>> batch_taglists;

        for (size_t i = 0; i < batch_count_; ++i) {
          batch_sparse_count.push_back(sparse_data[idx].count);

          for (size_t k = 0; k < sparse_data[idx].indices.size(); ++k) {
            batch_sparse_indices.push_back(sparse_data[idx].indices[k]);
          }

          for (size_t k = 0; k < sparse_data[idx].features.size(); ++k) {
            batch_sparse_feature.push_back(sparse_data[idx].features[k]);
          }

          if (taglists.size() > idx) {
            batch_taglists.push_back(taglists[idx]);
          }

          idx = (idx + 1) % queries.size();
        }

        batch_sparse_counts_.push_back(batch_sparse_count);
        batch_sparse_indices_.push_back(batch_sparse_indices);
        batch_sparse_features_.push_back(batch_sparse_feature);
        batch_taglists_.push_back(batch_taglists);
      }
    }

    if (typeid(T) == typeid(float)) {
      qmeta_.set_data_type(IndexMeta::DataType::DT_FP32);
    } else if (typeid(T) == typeid(int8_t)) {
      qmeta_.set_data_type(IndexMeta::DataType::DT_INT8);
    } else {
      LOG_ERROR("unsupported type");
      return false;
    }

    cout << "Load query done!" << endl;
    return true;
  }

  void run(core_interface::Index::Pointer index,
           core_interface::BaseIndexQueryParam::Pointer query_param,
           int max_iter, int topk) {
    // Check
    if (batch_sparse_counts_.size() == 0) {
      return;
    }

    query_param_ = query_param;
    query_param_->topk = topk;
    query_param_->is_linear = false;

    // Do bench
    signal(SIGINT, stop);
    bench_result_.mark_start();
    auto start_time = Monotime::MilliSeconds();
    for (size_t i = 0; i < threads_; ++i) {
      pool_->execute(this, &SparseBench<T>::start_bench, index, max_iter,
                     &STOP_NOW);
    }

    while (!pool_->is_finished()) {
      this_thread::sleep_for(chrono::milliseconds(1));
      if (Monotime::MilliSeconds() - start_time > bench_secs_ * 1000) {
        STOP_NOW = true;
      }
    }

    pool_->wait_finish();

    bench_result_.mark_end();
    bench_result_.print();
  }

  void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
                     const std::vector<uint64_t> &tag_key_list) {
    id_to_tags_list_ = id_to_tags_list;
    tag_key_list_ = tag_key_list;
  }

 private:
  void start_bench(core_interface::Index::Pointer index, size_t max_iter,
                   const bool *is_stop) {
    size_t thread_index = pool_->indexof_this();

    size_t i = thread_index;
    size_t sparse_query_size = batch_sparse_indices_.size();
    for (; i < max_iter && !*is_stop; i += threads_) {
      int idx = i % sparse_query_size;

      // prefilter
      FilterResultCache filter_cache;
      std::shared_ptr<IndexFilter> filter_ptr = nullptr;
      if (filter_mode_ == FM_TAG) {
        if (batch_taglists_[idx].size() != 1) {
          LOG_ERROR("query tag list not equal to one!");
          return;
        }

        int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
                                      tag_key_list_);
        if (ret != 0) {
          LOG_ERROR("prefilter failed, idx: %d", idx);
          return;
        }

        auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };

        filter_ptr = std::make_shared<IndexFilter>();
        filter_ptr->set(filterFunc);
      }

      auto query_param = query_param_->Clone();
      query_param->filter = filter_ptr;

      // Do knn_search
      uint64_t start = Monotime::MicroSeconds();
      int ret;
      if (batch_count_ == 1) {
        if (batch_sparse_counts_[idx].size() != 1) {
          LOG_ERROR("Sparse count size should be 1, since batch count is 1");
          return;
        }
        ret = do_knn_search<T>(index, batch_sparse_counts_[idx][0],
                               batch_sparse_indices_[idx],
                               batch_sparse_features_[idx], query_param);
      } else {
        ret = do_knn_search_batch<T>(
            index, batch_sparse_counts_[idx], batch_sparse_indices_[idx],
            batch_sparse_features_[idx], idx, query_param);
      }

      if (ret != 0) {
        LOG_ERROR("Failed to sparse knn search, ret=%d %s", ret,
                  IndexError::What(ret));
        return;
      }

      uint64_t end = Monotime::MicroSeconds();

      // Do sample
      bench_result_.add_time(batch_count_, end - start);
    }
  }

  // sparse search - single query
  template <typename U>
  typename std::enable_if<std::is_same<float, U>::value, int>::type
  do_knn_search(core_interface::Index::Pointer index,
                const uint32_t sparse_count,
                const vector<uint32_t> &sparse_indices,
                const vector<U> &sparse_feature,
                core_interface::BaseIndexQueryParam::Pointer query_param) {
    core_interface::SparseVector sparse_query;
    sparse_query.count = sparse_count;
    sparse_query.indices = sparse_indices.data();
    sparse_query.values = sparse_feature.data();
    core_interface::VectorData query_data;
    query_data.vector = sparse_query;

    core_interface::SearchResult search_result;
    int ret = index->Search(query_data, query_param, &search_result);
    if (ret < 0) {
      return ret;
    }

    if (search_result.doc_list_.empty()) {
      LOG_ERROR("Search results is empty");
    }

    return 0;
  }

  template <typename U>
  typename std::enable_if<std::is_same<int8_t, U>::value ||
                              std::is_same<uint32_t, U>::value ||
                              std::is_same<uint64_t, U>::value,
                          int>::type
  do_knn_search(core_interface::Index::Pointer /*index*/,
                const uint32_t /*sparse_count*/,
                const vector<uint32_t> & /*sparse_indices*/,
                const vector<U> & /*sparse_feature*/,
                core_interface::BaseIndexQueryParam::Pointer /*query_param*/) {
    return IndexError_Unsupported;
  }

  // sparse search - batch
  template <typename U>
  typename std::enable_if<std::is_same<float, U>::value, int>::type
  do_knn_search_batch(
      core_interface::Index::Pointer index,
      const vector<uint32_t> &sparse_count,
      const vector<uint32_t> & /*sparse_indices*/,
      const vector<U> & /*sparse_feature*/, size_t batch_idx,
      core_interface::BaseIndexQueryParam::Pointer query_param) {
    // For batch search, search each query separately
    for (size_t i = 0; i < sparse_count.size(); ++i) {
      size_t query_idx = batch_idx * batch_count_ + i;
      if (query_idx >= linear_sparse_data_.size()) {
        break;
      }

      const auto &single_sparse = linear_sparse_data_[query_idx];
      core_interface::SparseVector sparse_query;
      sparse_query.count = single_sparse.count;
      sparse_query.indices = single_sparse.indices.data();
      sparse_query.values = single_sparse.features.data();
      core_interface::VectorData query_data;
      query_data.vector = sparse_query;

      core_interface::SearchResult search_result;
      int ret = index->Search(query_data, query_param, &search_result);
      if (ret < 0) {
        return ret;
      }

      if (search_result.doc_list_.empty()) {
        LOG_ERROR("Search results is empty for batch query %zu", i);
      }
    }

    return 0;
  }

  template <typename U>
  typename std::enable_if<std::is_same<int8_t, U>::value ||
                              std::is_same<uint32_t, U>::value ||
                              std::is_same<uint64_t, U>::value,
                          int>::type
  do_knn_search_batch(
      core_interface::Index::Pointer /*index*/,
      const vector<uint32_t> & /*sparse_count*/,
      const vector<uint32_t> & /*sparse_indices*/,
      const vector<U> & /*sparse_feature*/, size_t /*batch_idx*/,
      core_interface::BaseIndexQueryParam::Pointer /*query_param*/) {
    return IndexError_Unsupported;
  }

 private:
  IndexQueryMeta qmeta_{};
  size_t threads_;
  size_t bench_secs_;
  size_t batch_count_;
  core_interface::BaseIndexQueryParam::Pointer query_param_;
  shared_ptr<ThreadPool> pool_;

  vector<SparseData<T>> linear_sparse_data_;
  vector<vector<uint32_t>> batch_sparse_counts_;
  vector<vector<uint32_t>> batch_sparse_indices_;
  vector<vector<T>> batch_sparse_features_;
  vector<vector<vector<uint64_t>>> batch_taglists_;

  // Tag lists for filtering
  std::vector<std::vector<uint64_t>> id_to_tags_list_;
  std::vector<uint64_t> tag_key_list_;

  FilterMode filter_mode_{FM_NONE};
  BenchResult bench_result_;
  static bool STOP_NOW;
};
template <typename T>
bool SparseBench<T>::STOP_NOW = false;

bool check_config(YAML::Node &config_node) {
  auto common = config_node["IndexCommon"];
  if (!common) {
    LOG_ERROR("Can not find [IndexCommon] in config");
    return false;
  }
  if (!common["IndexConfig"]) {
    LOG_ERROR("Can not find [IndexConfig] in config");
    return false;
  }
  if (!common["IndexPath"]) {
    LOG_ERROR("Can not find [IndexPath] in config");
    return false;
  }
  if (!common["TopK"]) {
    LOG_ERROR("Can not find [TopK] in config");
    return false;
  }
  if (!common["QueryFile"]) {
    LOG_ERROR("Can not find [QueryFile] in config");
    return false;
  }

  auto query_config = config_node["QueryConfig"];
  if (!query_config) {
    LOG_ERROR("Can not find [QueryConfig] in config");
    return false;
  }
  if (!query_config["QueryParam"]) {
    LOG_ERROR("Can not find [QueryConfig.QueryParam] in config");
    return false;
  }


  return true;
}

void usage(void) {
  cout << "Usage: bench CONFIG.yaml [plugin file path]" << endl;
}


int bench(std::string &query_type, size_t thread_count, size_t batch_count,
          size_t top_k, string query_file, string &first_sep,
          string &second_sep, size_t bench_secs, size_t iter_count,
          core_interface::Index::Pointer index,
          core_interface::BaseIndexQueryParam::Pointer query_param,
          string &index_dir, RetrievalMode retrieval_mode,
          FilterMode filter_mode) {
  if (filter_mode == FM_TAG && batch_count > 1) {
    LOG_ERROR("filter mode can not be run in batch mode");
    return -1;
  }

  std::vector<std::vector<uint64_t>> id_to_tags_list;
  std::vector<uint64_t> tag_key_list;
  // Load tag lists if available
  load_taglists(index_dir, id_to_tags_list, tag_key_list);

  if (query_type == "float") {
    Bench<float> bench(thread_count, bench_secs, batch_count, retrieval_mode,
                       filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else if (query_type == "int8") {
    Bench<int8_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
                        filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else if (query_type == "binary") {
    Bench<uint32_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
                          filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else if (query_type == "binary64") {
    Bench<uint64_t> bench(thread_count, bench_secs, batch_count, retrieval_mode,
                          filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else {
    LOG_ERROR("Can not recognize type: %s", query_type.c_str());
  }

  return 0;
}

int bench_sparse(std::string &query_type, size_t thread_count,
                 size_t batch_count, size_t top_k, string query_file,
                 string &first_sep, string &second_sep, size_t bench_secs,
                 size_t iter_count, core_interface::Index::Pointer index,
                 core_interface::BaseIndexQueryParam::Pointer query_param,
                 string &index_dir, FilterMode filter_mode) {
  if (filter_mode == FM_TAG && batch_count > 1) {
    LOG_ERROR("filter mode can not be run in batch mode");
    return -1;
  }

  std::vector<std::vector<uint64_t>> id_to_tags_list;
  std::vector<uint64_t> tag_key_list;
  // Load tag lists if available
  load_taglists(index_dir, id_to_tags_list, tag_key_list);

  if (query_type == "float") {
    SparseBench<float> bench(thread_count, bench_secs, batch_count,
                             filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else if (query_type == "int8") {
    SparseBench<int8_t> bench(thread_count, bench_secs, batch_count,
                              filter_mode);
    bench.load_query(query_file, first_sep, second_sep);
    bench.set_tag_lists(id_to_tags_list, tag_key_list);
    bench.run(index, query_param, iter_count, top_k);
  } else {
    LOG_ERROR("Can not recognize type: %s", query_type.c_str());
  }

  return 0;
}

int main(int argc, char *argv[]) {
  if (argc < 2) {
    usage();
    return -1;
  }

  IndexPluginBroker broker;
  std::string error;
  for (int i = 2; i < argc; ++i) {
    if (!broker.emplace(argv[i], &error)) {
      LOG_ERROR("Failed to load plugin: %s (%s)", argv[i], error.c_str());
      return -1;
    }
  }

  YAML::Node config_node;
  try {
    config_node = YAML::LoadFile(argv[1]);
  } catch (...) {
    LOG_ERROR("Load YAML file[%s] failed!", argv[1]);
    return -1;
  }

  if (!check_config(config_node)) {
    return -1;
  }
  auto config_common = config_node["IndexCommon"];

  map<string, int> LOG_LEVEL = {{"debug", IndexLogger::LEVEL_DEBUG},
                                {"info", IndexLogger::LEVEL_INFO},
                                {"warn", IndexLogger::LEVEL_WARN},
                                {"error", IndexLogger::LEVEL_ERROR},
                                {"fatal", IndexLogger::LEVEL_FATAL}};
  string log_level = config_common["LogLevel"]
                         ? config_common["LogLevel"].as<string>()
                         : "debug";
  transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower);
  if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) {
    IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]);
    zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]);
  }

  // Calculate Bench
  size_t thread_count = config_common["BenchThreadCount"]
                            ? config_common["BenchThreadCount"].as<uint64_t>()
                            : 0;
  size_t iter_count = config_common["BenchIterCount"]
                          ? config_common["BenchIterCount"].as<uint64_t>()
                          : 10000;
  size_t batch_count = config_common["BenchBatchCount"]
                           ? config_common["BenchBatchCount"].as<uint64_t>()
                           : 0;
  g_debug_mode = config_common["DebugMode"]
                     ? config_common["DebugMode"].as<bool>()
                     : false;
  string topk_str = config_common["TopK"].as<string>();

  RetrievalMode retrieval_mode{RM_DENSE};
  if (config_common["RetrievalMode"]) {
    std::string retrieval_mode_str =
        config_common["RetrievalMode"].as<string>();
    if (retrieval_mode_str == "dense") {
      retrieval_mode = RM_DENSE;
    } else if (retrieval_mode_str == "sparse") {
      retrieval_mode = RM_SPARSE;
    }
  }

  FilterMode filter_mode{FM_NONE};
  if (config_common["FilterMode"]) {
    std::string filter_mode_str = config_common["FilterMode"].as<string>();
    if (filter_mode_str == "tag") {
      filter_mode = FM_TAG;
    }
  }

  vector<int32_t> topk_values;
  StringHelper::Split(topk_str, ",", &topk_values);
  size_t top_k = *topk_values.rbegin();
  string query_file = config_common["QueryFile"].as<string>();
  string first_sep = config_common["QueryFirstSep"]
                         ? config_common["QueryFirstSep"].as<string>()
                         : ";";
  string second_sep = config_common["QuerySecondSep"]
                          ? config_common["QuerySecondSep"].as<string>()
                          : " ";
  string query_type = config_common["QueryType"]
                          ? config_common["QueryType"].as<string>()
                          : "float";
  size_t bench_secs = config_common["BenchSecs"]
                          ? config_common["BenchSecs"].as<uint64_t>()
                          : 60;

  string index_dir = config_common["IndexPath"].as<string>();

  core_interface::Index::Pointer index;
  core_interface::BaseIndexQueryParam::Pointer query_param;
  if (0 !=
      parse_and_load_index_param(config_node, index_dir, index, query_param)) {
    LOG_ERROR("Failed to parse and load index param");
    return -1;
  }

  if (retrieval_mode == RM_SPARSE) {
    bench_sparse(query_type, thread_count, batch_count, top_k, query_file,
                 first_sep, second_sep, bench_secs, iter_count, index,
                 query_param, index_dir, filter_mode);

    cout << "Bench Sparse done." << endl;
  } else {
    bench(query_type, thread_count, batch_count, top_k, query_file, first_sep,
          second_sep, bench_secs, iter_count, index, query_param, index_dir,
          retrieval_mode, filter_mode);

    cout << "Bench done." << endl;
  }

  // Cleanup
  index->Close();

  return 0;
}