/*
 * Copyright (c) 2025, valkey-search contributors
 * All rights reserved.
 * SPDX-License-Identifier: BSD 3-Clause
 *
 */

#include "src/query/search.h"

#include <absl/strings/str_split.h>

#include <cstddef>
#include <deque>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "src/attribute_data_type.h"
#include "src/indexes/index_base.h"
#include "src/indexes/numeric.h"
#include "src/indexes/tag.h"
#include "src/indexes/text.h"
#include "src/indexes/text/orproximity.h"
#include "src/indexes/text/proximity.h"
#include "src/indexes/text/text_fetcher.h"
#include "src/indexes/universal_set_fetcher.h"
#include "src/indexes/vector_base.h"
#include "src/indexes/vector_flat.h"
#include "src/indexes/vector_hnsw.h"
#include "src/metrics.h"
#include "src/query/content_resolution.h"
#include "src/query/planner.h"
#include "src/query/predicate.h"
#include "src/valkey_search.h"
#include "src/valkey_search_options.h"
#include "third_party/hnswlib/hnswlib.h"
#include "vmsdk/src/debug.h"
#include "vmsdk/src/info.h"
#include "vmsdk/src/latency_sampler.h"
#include "vmsdk/src/log.h"
#include "vmsdk/src/managed_pointers.h"
#include "vmsdk/src/status/status_macros.h"
#include "vmsdk/src/thread_pool.h"
#include "vmsdk/src/time_sliced_mrmw_mutex.h"
#include "vmsdk/src/type_conversions.h"
#include "vmsdk/src/valkey_module_api/valkey_module.h"

namespace valkey_search::query {

// Query operation counters
DEV_INTEGER_COUNTER(query_stats, query_text_term_count);
DEV_INTEGER_COUNTER(query_stats, query_text_prefix_count);
DEV_INTEGER_COUNTER(query_stats, query_text_suffix_count);
DEV_INTEGER_COUNTER(query_stats, query_text_fuzzy_count);
DEV_INTEGER_COUNTER(query_stats, query_text_proximity_count);
DEV_INTEGER_COUNTER(query_stats, query_numeric_count);
DEV_INTEGER_COUNTER(query_stats, query_tag_count);
DEV_INTEGER_COUNTER(query_stats, nonvector_results_fetched_limited_count);

class InlineVectorFilter : public hnswlib::BaseFilterFunctor {
 public:
  InlineVectorFilter(
      query::Predicate *filter_predicate, indexes::VectorBase *vector_index,
      const std::shared_ptr<indexes::text::TextIndexSchema> text_index_schema,
      QueryOperations query_operations)
      : filter_predicate_(filter_predicate),
        vector_index_(vector_index),
        text_index_schema_(text_index_schema),
        query_operations_(query_operations) {}
  ~InlineVectorFilter() override = default;

  bool operator()(hnswlib::labeltype id) override {
    BACKGROUND_PAUSEPOINT("search_inline_filter");
    auto key = vector_index_->GetKeyDuringSearch(id);
    if (!key.ok()) {
      return false;
    }
    const valkey_search::indexes::text::TextIndex *text_index = nullptr;
    if (text_index_schema_) {
      text_index = text_index_schema_->GetPerKeyTextIndex(*key, false);
    }
    indexes::PrefilterEvaluator evaluator(text_index, query_operations_);
    return evaluator.Evaluate(*filter_predicate_, *key);
  }

 private:
  query::Predicate *filter_predicate_;
  indexes::VectorBase *vector_index_;
  const std::shared_ptr<indexes::text::TextIndexSchema> text_index_schema_;
  QueryOperations query_operations_;
};
absl::StatusOr<std::vector<indexes::Neighbor>> PerformVectorSearch(
    indexes::VectorBase *vector_index, const SearchParameters &parameters) {
  std::unique_ptr<InlineVectorFilter> inline_filter;
  if (parameters.filter_parse_results.root_predicate != nullptr) {
    const std::shared_ptr<indexes::text::TextIndexSchema> text_index_schema =
        parameters.index_schema->GetTextIndexSchema();
    inline_filter = std::make_unique<InlineVectorFilter>(
        parameters.filter_parse_results.root_predicate.get(), vector_index,
        text_index_schema, parameters.filter_parse_results.query_operations);
    VMSDK_LOG(DEBUG, nullptr) << "Performing vector search with inline filter";
  }
  if (vector_index->GetIndexerType() == indexes::IndexerType::kHNSW) {
    auto vector_hnsw = dynamic_cast<indexes::VectorHNSW<float> *>(vector_index);

    auto latency_sample = SAMPLE_EVERY_N(100);
    auto res = vector_hnsw->Search(parameters.query, parameters.k,
                                   parameters.cancellation_token,
                                   std::move(inline_filter), parameters.ef,
                                   parameters.enable_partial_results);
    Metrics::GetStats().hnsw_vector_index_search_latency.SubmitSample(
        std::move(latency_sample));
    return res;
  }
  if (vector_index->GetIndexerType() == indexes::IndexerType::kFlat) {
    auto vector_flat = dynamic_cast<indexes::VectorFlat<float> *>(vector_index);
    auto latency_sample = SAMPLE_EVERY_N(100);
    auto res = vector_flat->Search(parameters.query, parameters.k,
                                   parameters.cancellation_token,
                                   std::move(inline_filter));
    Metrics::GetStats().flat_vector_index_search_latency.SubmitSample(
        std::move(latency_sample));
    return res;
  }
  CHECK(false) << "Unsupported indexer type: "
               << (int)vector_index->GetIndexerType();
}

void AppendQueue(
    std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> &dest,
    std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> &src) {
  while (!src.empty()) {
    dest.push(std::move(src.front()));
    src.pop();
  }
}

inline PredicateType EvaluateAsComposedPredicate(
    const Predicate *composed_predicate, bool negate) {
  auto predicate_type = composed_predicate->GetType();

  if (!negate) {
    return predicate_type;
  }
  if (predicate_type == PredicateType::kComposedAnd) {
    return PredicateType::kComposedOr;
  }
  return PredicateType::kComposedAnd;
}

// Helper fn to identify if query is not fully solved after the entries fetcher
// search, meaning it requires prefilter evaluation Prefiltering is needed when
// query contains an AND with numeric or tag predicates.
// It is also needed when negate is involved.
inline bool IsUnsolvedQuery(QueryOperations query_operations) {
  return query_operations & (QueryOperations::kContainsNumeric |
                             QueryOperations::kContainsTag) &&
             query_operations & QueryOperations::kContainsAnd ||
         (query_operations & QueryOperations::kContainsNegate);
}

// Helper fn to identify if deduplication is needed.
// (1) OR operations need deduplication.
// (2) Any TAG operations need deduplication.
// (3) Non-text negation needs deduplication (uses NegateEntriesFetcher)
inline bool NeedsDeduplication(QueryOperations query_operations) {
  bool has_or = query_operations & QueryOperations::kContainsOr;
  bool has_tag = query_operations & QueryOperations::kContainsTag;
  bool has_negate = query_operations & QueryOperations::kContainsNegate;
  bool has_text = query_operations & QueryOperations::kContainsText;
  // Text + negate doesn't need dedup (handled by prefilter evaluation)
  if (has_text && has_negate) {
    return false;
  }
  return has_or || has_tag || has_negate;
}

// Builds TextIterator for text predicates. Returns pair of iterator and
// estimated size.
std::pair<std::unique_ptr<indexes::text::TextIterator>, size_t>
BuildTextIterator(const Predicate *predicate, bool negate,
                  bool require_positions, bool is_vec_query) {
  if (predicate->GetType() == PredicateType::kComposedAnd ||
      predicate->GetType() == PredicateType::kComposedOr) {
    auto composed_predicate =
        dynamic_cast<const ComposedPredicate *>(predicate);
    auto predicate_type =
        EvaluateAsComposedPredicate(composed_predicate, negate);
    auto slop = composed_predicate->GetSlop();
    bool inorder = composed_predicate->GetInorder();
    bool child_require_positions = slop.has_value() || inorder;
    if (predicate_type == PredicateType::kComposedAnd) {
      absl::InlinedVector<std::unique_ptr<indexes::text::TextIterator>,
                          indexes::text::kProximityTermsInlineCapacity>
          iterators;
      size_t min_size = SIZE_MAX;
      for (const auto &child : composed_predicate->GetChildren()) {
        auto [iter, size] = BuildTextIterator(
            child.get(), negate, child_require_positions, is_vec_query);
        if (iter) {
          iterators.push_back(std::move(iter));
          min_size = std::min(min_size, size);
        }
      }
      // The Composed AND only has non text predicates, return null
      // to have the caller handle it.
      if (iterators.empty()) return {nullptr, 0};
      bool skip_positional = !child_require_positions;
      size_t total_size = min_size == SIZE_MAX ? 0 : min_size;
      return {std::make_unique<indexes::text::ProximityIterator>(
                  std::move(iterators), slop, inorder, skip_positional),
              total_size};
    } else {
      absl::InlinedVector<std::unique_ptr<indexes::text::TextIterator>,
                          indexes::text::kProximityTermsInlineCapacity>
          iterators;
      size_t total_size = 0;
      bool has_non_text = false;
      for (const auto &child : composed_predicate->GetChildren()) {
        auto [iter, size] = BuildTextIterator(
            child.get(), negate, child_require_positions, is_vec_query);
        if (iter) {
          iterators.push_back(std::move(iter));
          total_size += size;
        } else {
          has_non_text = true;
        }
      }
      // If the Composed OR has any non text predicate, we cannot
      // build a text iterator.
      if (iterators.empty() || has_non_text) return {nullptr, 0};
      return {std::make_unique<indexes::text::OrProximityIterator>(
                  std::move(iterators)),
              total_size};
    }
  }
  if (predicate->GetType() == PredicateType::kText) {
    auto text_predicate = dynamic_cast<const TextPredicate *>(predicate);
    auto text_index = text_predicate->GetTextIndexSchema()->GetTextIndex();
    auto field_mask = text_predicate->GetFieldMask();
    size_t size = text_predicate->EstimateSize(is_vec_query);
    auto result = text_predicate->BuildTextIterator(text_index, field_mask,
                                                    require_positions);
    return {std::move(result), size};
  }
  if (predicate->GetType() == PredicateType::kNegate) {
    // Cannot build text iterator for negation - return null
    return {nullptr, 0};
  }
  // Numeric/Tag
  return {nullptr, 0};
}

size_t EvaluateFilterAsPrimary(
    const SearchParameters &parameters, const Predicate *predicate,
    std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> &entries_fetchers,
    bool negate) {
  const QueryOperations query_operations =
      parameters.filter_parse_results.query_operations;
  const IndexSchema *index_schema = parameters.index_schema.get();
  const bool is_vec_query = parameters.IsVectorQuery();

  // Always use universal set when query has text + negate
  if ((query_operations & QueryOperations::kContainsText) &&
      (query_operations & QueryOperations::kContainsNegate)) {
    CHECK(index_schema != nullptr) << "IndexSchema required for text+negate";
    auto universal_fetcher =
        std::make_unique<indexes::UniversalSetFetcher>(index_schema);
    size_t size = universal_fetcher->Size();
    entries_fetchers.push(std::move(universal_fetcher));
    return size;
  }

  if (predicate->GetType() == PredicateType::kComposedAnd ||
      predicate->GetType() == PredicateType::kComposedOr) {
    auto composed_predicate =
        dynamic_cast<const ComposedPredicate *>(predicate);
    auto predicate_type =
        EvaluateAsComposedPredicate(composed_predicate, negate);
    if (predicate_type == PredicateType::kComposedAnd) {
      auto [text_iter, size] =
          BuildTextIterator(composed_predicate, negate, false, is_vec_query);
      if (text_iter) {
        entries_fetchers.push(
            std::make_unique<indexes::text::TextIteratorFetcher>(
                std::move(text_iter), size));
        return size;
      }
      size_t min_size = SIZE_MAX;
      std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> best_fetchers;
      for (const auto &child : composed_predicate->GetChildren()) {
        std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> child_fetchers;
        size_t child_size = EvaluateFilterAsPrimary(parameters, child.get(),
                                                    child_fetchers, negate);
        if (child_size < min_size) {
          min_size = child_size;
          best_fetchers = std::move(child_fetchers);
        }
      }
      AppendQueue(entries_fetchers, best_fetchers);
      return min_size;
    } else {
      size_t total_size = 0;
      for (const auto &child : composed_predicate->GetChildren()) {
        std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> child_fetchers;
        size_t child_size = EvaluateFilterAsPrimary(parameters, child.get(),
                                                    child_fetchers, negate);
        AppendQueue(entries_fetchers, child_fetchers);
        total_size += child_size;
      }
      return total_size;
    }
  }
  if (predicate->GetType() == PredicateType::kTag) {
    auto tag_predicate = dynamic_cast<const TagPredicate *>(predicate);
    auto fetcher = tag_predicate->GetIndex()->Search(*tag_predicate, negate);
    size_t size = fetcher->Size();
    entries_fetchers.push(std::move(fetcher));
    return size;
  }
  if (predicate->GetType() == PredicateType::kNumeric) {
    auto numeric_predicate = dynamic_cast<const NumericPredicate *>(predicate);
    auto fetcher =
        numeric_predicate->GetIndex()->Search(*numeric_predicate, negate);
    size_t size = fetcher->Size();
    entries_fetchers.push(std::move(fetcher));
    return size;
  }
  if (predicate->GetType() == PredicateType::kText) {
    auto text_predicate = dynamic_cast<const TextPredicate *>(predicate);
    size_t size = text_predicate->EstimateSize(is_vec_query);
    auto fetcher = std::make_unique<indexes::Text::EntriesFetcher>(
        size, text_predicate->GetTextIndexSchema()->GetTextIndex(),
        text_predicate->GetFieldMask(), false);
    fetcher->predicate_ = text_predicate;
    entries_fetchers.push(std::move(fetcher));
    return size;
  }
  if (predicate->GetType() == PredicateType::kNegate) {
    auto negate_predicate = dynamic_cast<const NegatePredicate *>(predicate);
    size_t result =
        EvaluateFilterAsPrimary(parameters, negate_predicate->GetPredicate(),
                                entries_fetchers, !negate);
    return result;
  }
  CHECK(false);
}

struct PrefilteredKey {
  std::string key;
  float distance;
};

void EvaluatePrefilteredKeys(
    const SearchParameters &parameters,
    std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> &entries_fetchers,
    absl::AnyInvocable<bool(const InternedStringPtr &,
                            absl::flat_hash_set<const char *> &)>
        appender,
    size_t max_keys, bool stop_on_fetch_limit) {
  // If there was a union operation, we need to handle deduplication.
  // This implementation skips deduplication (flat_hash_set usage) if not needed
  // for performance.
  bool needs_dedup =
      NeedsDeduplication(parameters.filter_parse_results.query_operations);
  absl::flat_hash_set<const char *> result_keys;
  if (needs_dedup) {
    result_keys.reserve(max_keys);
  }
  const std::shared_ptr<indexes::text::TextIndexSchema> text_index_schema =
      parameters.index_schema ? parameters.index_schema->GetTextIndexSchema()
                              : nullptr;
  while (!entries_fetchers.empty()) {
    auto fetcher = std::move(entries_fetchers.front());
    entries_fetchers.pop();
    auto iterator = fetcher->Begin();
    while (!iterator->Done()) {
      const auto &key = **iterator;
      // 1. Skip if already processed (only if dedup is needed)
      if (needs_dedup && result_keys.contains(key->Str().data())) {
        iterator->Next();
        continue;
      }
      const valkey_search::indexes::text::TextIndex *text_index =
          text_index_schema ? text_index_schema->GetPerKeyTextIndex(key, false)
                            : nullptr;
      indexes::PrefilterEvaluator key_evaluator(
          text_index, parameters.filter_parse_results.query_operations);
      BACKGROUND_PAUSEPOINT("search_prefilter_eval");
      // 3. Evaluate predicate
      if (key_evaluator.Evaluate(
              *parameters.filter_parse_results.root_predicate, key)) {
        bool result = appender(key, result_keys);
        if (needs_dedup && result) {
          result_keys.insert(key->Str().data());
        }
        // For non-vector queries that exceed the fetch limit, return early
        if (stop_on_fetch_limit && !result) {
          return;
        }
      }
      iterator->Next();
      if (parameters.cancellation_token->IsCancelled()) {
        return;
      }
    }
  }
}

std::priority_queue<std::pair<float, hnswlib::labeltype>>
CalcBestMatchingPrefilteredKeys(
    const SearchParameters &parameters,
    std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> &entries_fetchers,
    indexes::VectorBase *vector_index, size_t qualified_entries) {
  std::priority_queue<std::pair<float, hnswlib::labeltype>> results;
  auto results_appender =
      [&results, &parameters, vector_index](
          const InternedStringPtr &key,
          absl::flat_hash_set<const char *> &top_keys) -> bool {
    return vector_index->AddPrefilteredKey(parameters.query, parameters.k, key,
                                           results, top_keys);
  };
  EvaluatePrefilteredKeys(parameters, entries_fetchers,
                          std::move(results_appender), qualified_entries,
                          /*stop_on_fetch_limit=*/false);
  return results;
}

std::string StringFormatVector(std::vector<char> vector) {
  if (vector.size() % sizeof(float) != 0) {
    return {vector.data(), vector.size()};
  }

  std::vector<std::string> float_strings;
  for (size_t i = 0; i < vector.size(); i += sizeof(float)) {
    float value;
    std::memcpy(&value, vector.data() + i, sizeof(float));
    float_strings.push_back(absl::StrCat(value));
  }

  return absl::StrCat("[", absl::StrJoin(float_strings, ","), "]");
}

absl::StatusOr<std::vector<indexes::Neighbor>> MaybeAddIndexedContent(
    absl::StatusOr<std::vector<indexes::Neighbor>> results,
    const SearchParameters &parameters) {
  if (!results.ok()) {
    return results;
  }
  if (parameters.no_content || parameters.return_attributes.empty()) {
    return results;
  }
  struct AttributeInfo {
    const ReturnAttribute *attribute;
    indexes::IndexBase *index;
  };
  std::vector<AttributeInfo> attributes;
  for (auto &attribute : parameters.return_attributes) {
    if (!attribute.attribute_alias.get()) {
      // Any attribute that is not indexed will result in all attributes being
      // fetched from the main thread for consistency.
      return results;
    }
    auto index = parameters.index_schema->GetIndex(
        vmsdk::ToStringView(attribute.attribute_alias.get()));
    if (!index.ok()) {
      return results;
    }
    attributes.push_back(AttributeInfo{&attribute, index.value().get()});
  }
  for (auto &neighbor : *results) {
    if (neighbor.attribute_contents.has_value()) {
      continue;
    }
    neighbor.attribute_contents = RecordsMap();
    bool any_value_missing = false;
    for (auto &attribute_info : attributes) {
      vmsdk::UniqueValkeyString attribute_value = nullptr;
      switch (attribute_info.index->GetIndexerType()) {
        case indexes::IndexerType::kTag: {
          auto tag_index = dynamic_cast<indexes::Tag *>(attribute_info.index);
          auto tag_value_ptr = tag_index->GetRawValue(neighbor.external_id);
          if (tag_value_ptr) {
            attribute_value = vmsdk::MakeUniqueValkeyString(*tag_value_ptr);
          }
          break;
        }
        case indexes::IndexerType::kNumeric: {
          auto numeric_index =
              dynamic_cast<indexes::Numeric *>(attribute_info.index);
          auto numeric = numeric_index->GetValue(neighbor.external_id);
          if (numeric != nullptr) {
            attribute_value =
                vmsdk::MakeUniqueValkeyString(absl::StrCat(*numeric));
          }
          break;
        }
        case indexes::IndexerType::kVector:
        case indexes::IndexerType::kHNSW:
        case indexes::IndexerType::kFlat: {
          auto vector_index =
              dynamic_cast<indexes::VectorBase *>(attribute_info.index);
          auto vector = vector_index->GetValue(neighbor.external_id);
          if (vector.ok()) {
            if (parameters.index_schema->GetAttributeDataType().ToProto() ==
                data_model::AttributeDataType::ATTRIBUTE_DATA_TYPE_JSON) {
              attribute_value = vmsdk::MakeUniqueValkeyString(
                  StringFormatVector(vector.value()));
            } else {
              attribute_value =
                  vmsdk::UniqueValkeyString(ValkeyModule_CreateString(
                      nullptr, vector->data(), vector->size()));
            }
          } else {
            VMSDK_LOG_EVERY_N_SEC(WARNING, nullptr, 1)
                << "Failed to get vector value during fetching through index "
                   "contents: "
                << vector.status();
          }
          break;
        }
        case indexes::IndexerType::kText: {
          // Text indexes don't store retrievable raw values
          any_value_missing = true;
          break;
        }
        default:
          CHECK(false) << "Unsupported indexer type: "
                       << (int)attribute_info.index->GetIndexerType();
      }

      if (attribute_value != nullptr) {
        auto identifier = vmsdk::MakeUniqueValkeyString(
            vmsdk::ToStringView(attribute_info.attribute->identifier.get()));
        auto identifier_view = vmsdk::ToStringView(identifier.get());
        neighbor.attribute_contents->emplace(
            identifier_view,
            RecordsMapValue(std::move(identifier), std::move(attribute_value)));
      } else {
        // Mark this neighbor as needing content retrieval via the main thread
        // (e.g. the attribute value may exist but not be indexed due to type
        // mismatch).
        any_value_missing = true;
        break;
      }
    }
    if (any_value_missing) {
      neighbor.attribute_contents = std::nullopt;
    }
  }
  return results;
}

absl::StatusOr<std::vector<indexes::Neighbor>> SearchNonVectorQuery(
    const SearchParameters &parameters) {
  std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> entries_fetchers;
  size_t qualified_entries = EvaluateFilterAsPrimary(
      parameters, parameters.filter_parse_results.root_predicate.get(),
      entries_fetchers, false);

  // Get the config for maximum number of keys to accumulate before content
  // fetching
  const size_t max_keys = static_cast<size_t>(
      options::GetMaxNonVectorSearchResultsFetched().GetValue());
  std::vector<indexes::Neighbor> neighbors;
  neighbors.reserve(std::min(qualified_entries, static_cast<size_t>(5000)));
  bool fetch_limited = false;
  auto results_appender =
      [&neighbors, &parameters, max_keys, &fetch_limited](
          const InternedStringPtr &key,
          absl::flat_hash_set<const char *> &top_keys) -> bool {
    if (neighbors.size() >= max_keys) {
      fetch_limited = true;
      return false;
    }
    neighbors.emplace_back(indexes::Neighbor{key, 0.0f});
    return true;
  };
  // Cannot skip evaluation if the query contains unsolved composed operations.
  bool requires_prefilter_evaluation =
      IsUnsolvedQuery(parameters.filter_parse_results.query_operations);
  if (!requires_prefilter_evaluation) {
    bool needs_dedup =
        NeedsDeduplication(parameters.filter_parse_results.query_operations);
    absl::flat_hash_set<const char *> seen_keys;
    if (needs_dedup) {
      seen_keys.reserve(std::min(qualified_entries, static_cast<size_t>(5000)));
    }
    while (!entries_fetchers.empty()) {
      auto fetcher = std::move(entries_fetchers.front());
      entries_fetchers.pop();
      auto iterator = fetcher->Begin();
      while (!iterator->Done()) {
        const auto &key = **iterator;
        BACKGROUND_PAUSEPOINT("search_entries_fetcher");
        if (needs_dedup) {
          if (seen_keys.contains(key->Str().data())) {
            iterator->Next();
            continue;
          }
          seen_keys.insert(key->Str().data());
        }
        // Check if we've reached the limit
        if (neighbors.size() >= max_keys) {
          nonvector_results_fetched_limited_count.Increment();
          return neighbors;
        }
        neighbors.emplace_back(indexes::Neighbor{key, 0.0f});
        iterator->Next();
        if (parameters.cancellation_token->IsCancelled()) {
          return neighbors;
        }
      }
    }
    return neighbors;
  }
  EvaluatePrefilteredKeys(parameters, entries_fetchers,
                          std::move(results_appender), qualified_entries,
                          /*stop_on_fetch_limit=*/true);
  if (fetch_limited) {
    nonvector_results_fetched_limited_count.Increment();
  }
  return neighbors;
}

absl::StatusOr<std::vector<indexes::Neighbor>> DoSearch(
    const SearchParameters &parameters, SearchMode search_mode,
    vmsdk::ReaderMutexLock &lock) {
  ++Metrics::GetStats().time_slice_queries;
  // Handle OOM for search requests, defends against request
  // coming from the coordinator
  if (search_mode == SearchMode::kRemote) {
    auto ctx = vmsdk::MakeUniqueValkeyThreadSafeContext(nullptr);
    auto ctx_flags = ValkeyModule_GetContextFlags(ctx.get());
    if (ctx_flags & VALKEYMODULE_CTX_FLAGS_OOM) {
      return absl::ResourceExhaustedError(kOOMMsg);
    }
  }
  // Handle non vector queries first where attribute_alias is empty.
  if (parameters.IsNonVectorQuery()) {
    return SearchNonVectorQuery(parameters);
  }
  VMSDK_ASSIGN_OR_RETURN(auto index, parameters.index_schema->GetIndex(
                                         parameters.attribute_alias));
  auto vector_index = dynamic_cast<indexes::VectorBase *>(index.get());
  if (index->GetIndexerType() != indexes::IndexerType::kHNSW &&
      index->GetIndexerType() != indexes::IndexerType::kFlat) {
    return absl::InvalidArgumentError(
        absl::StrCat(parameters.attribute_alias, " is not a Vector index "));
  }

  if (!parameters.filter_parse_results.root_predicate) {
    return PerformVectorSearch(vector_index, parameters);
  }
  std::queue<std::unique_ptr<indexes::EntriesFetcherBase>> entries_fetchers;
  size_t qualified_entries = EvaluateFilterAsPrimary(
      parameters, parameters.filter_parse_results.root_predicate.get(),
      entries_fetchers, false);

  // Query planner makes the decision for pre-filtering vs inline-filtering.
  if (UsePreFiltering(qualified_entries, vector_index)) {
    VMSDK_LOG(DEBUG, nullptr)
        << "Using pre-filter query execution, qualified entries="
        << qualified_entries;
    // Do an exact nearest neighbour search on the reduced search space.
    ++Metrics::GetStats().query_prefiltering_requests_cnt;
    std::priority_queue<std::pair<float, hnswlib::labeltype>> results =
        CalcBestMatchingPrefilteredKeys(parameters, entries_fetchers,
                                        vector_index, qualified_entries);

    return vector_index->CreateReply(results);
  }
  ++Metrics::GetStats().query_inline_filtering_requests_cnt;
  lock.SetMayProlong();
  return PerformVectorSearch(vector_index, parameters);
}

// Check if no results should be returned based on query parameters.
// This handles two cases:
// 1. Any query with limit number == 0
// 2. Vector queries with limit first_index >= k
bool ShouldReturnNoResults(const SearchParameters &parameters) {
  return (parameters.IsVectorQuery() &&
          parameters.limit.first_index >=
              static_cast<uint64_t>(parameters.k)) ||
         parameters.limit.number == 0;
}

SearchResult::SearchResult()
    : total_count(0), is_limited_with_buffer(false), is_offsetted(false) {}

SearchResult::SearchResult(size_t total_count,
                           std::vector<indexes::Neighbor> neighbors,
                           const SearchParameters &parameters,
                           bool trim_offset_in_background)
    : total_count(total_count),
      is_limited_with_buffer(false),
      is_offsetted(false) {
  // Clear neighbors if no results should be returned
  if (ShouldReturnNoResults(parameters)) {
    this->neighbors.clear();
    return;
  }
  this->neighbors = std::move(neighbors);
  // Check if the command needs all results (e.g. for sorting). Trim otherwise.
  if (!parameters.RequiresCompleteResults()) {
    TrimResults(this->neighbors, parameters, trim_offset_in_background);
  }
}

// Apply limiting in background thread if possible.
void SearchResult::TrimResults(std::vector<indexes::Neighbor> &neighbors,
                               const SearchParameters &parameters,
                               bool trim_offset_in_background) {
  // Use the existing complex logic from SearchResult::GetSerializationRange
  SerializationRange range = GetSerializationRange(parameters);
  size_t max_needed = static_cast<size_t>(
      range.end_index * options::GetSearchResultBufferMultiplier());
  // In standalone mode, we can optimize by trimming from front first.
  // In cluster mode on remote searches on individual shards, we cannot trim
  // from the front yet because each shard produces X results. However, the
  // coordinator (after merging) WILL trim from the front and back in the
  // background thread to avoid memory bloat with large offsets / limit counts
  // before returning to the main thread.
  if (!ValkeySearch::Instance().IsCluster() || trim_offset_in_background) {
    this->is_offsetted = true;
    // Trim from front (apply offset)
    if (range.start_index > 0 && range.start_index < neighbors.size()) {
      neighbors.erase(neighbors.begin(), neighbors.begin() + range.start_index);
      // After trimming from the front, we no longer have an offset.
      // We only need (end_index - start_index) items.
      size_t actual_count = range.end_index - range.start_index;
      max_needed = static_cast<size_t>(
          actual_count * options::GetSearchResultBufferMultiplier());
    } else if (range.start_index >= neighbors.size()) {
      neighbors.clear();
      return;
    }
  }
  // If we don't need to limit, return early.
  if (neighbors.size() <= max_needed) {
    return;
  }
  // Apply limiting with buffer
  this->is_limited_with_buffer = true;
  neighbors.erase(neighbors.begin() + max_needed, neighbors.end());
}

// Determine the range of neighbors to serialize in the response.
SerializationRange SearchResult::GetSerializationRange(
    const SearchParameters &parameters) const {
  CHECK(!ShouldReturnNoResults(parameters));
  // Determine start_index
  size_t start_index = 0;
  // If we have already offsetted, start_index is 0.
  if (!is_offsetted) {
    if (parameters.IsVectorQuery()) {
      CHECK_GT(parameters.k, parameters.limit.first_index);
    }
    start_index = std::min(neighbors.size(),
                           static_cast<size_t>(parameters.limit.first_index));
  }
  // Determine end_index logic
  size_t limit_count = static_cast<size_t>(parameters.limit.number);
  size_t count;
  if (parameters.IsNonVectorQuery()) {
    count = std::min(limit_count, neighbors.size());
  } else {
    count = std::min(
        {static_cast<size_t>(parameters.k), limit_count, neighbors.size()});
  }
  size_t end_index = std::min(start_index + count, neighbors.size());
  // Return the range
  return {start_index, end_index};
}

absl::Status Search(SearchParameters &parameters, SearchMode search_mode) {
  auto &time_sliced_mutex = parameters.index_schema->GetTimeSlicedMutex();
  vmsdk::ReaderMutexLock lock(&time_sliced_mutex);
  absl::StatusOr<std::vector<indexes::Neighbor>> neighbors =
      DoSearch(parameters, search_mode, lock);
  VMSDK_ASSIGN_OR_RETURN(
      auto result, MaybeAddIndexedContent(std::move(neighbors), parameters));
  size_t total_count = result.size();
  parameters.search_result =
      SearchResult(total_count, std::move(result), parameters);
  parameters.index_schema->PopulateIndexMutationSequenceNumbers(
      parameters.search_result.neighbors);
  return absl::OkStatus();
}

absl::Status SearchAsync(std::unique_ptr<SearchParameters> parameters,
                         vmsdk::ThreadPool *thread_pool,
                         SearchMode search_mode) {
  thread_pool->Schedule(
      [parameters = std::move(parameters), search_mode]() mutable {
        auto res = Search(*parameters, search_mode);
        parameters->search_result.status = res;
        switch (parameters->GetContentProcessing()) {
          case ContentProcessing::kNoContent:
            parameters->QueryCompleteBackground(std::move(parameters));
            break;
          case ContentProcessing::kContentRequired:
          case ContentProcessing::kContentionCheckRequired:
            vmsdk::RunByMain([parameters = std::move(parameters)]() mutable {
              ResolveContent(std::move(parameters));
            });
            break;
          default:
            CHECK(false) << "Unknown content processing mode";
        }
      },
      vmsdk::ThreadPool::Priority::kHigh);
  return absl::OkStatus();
}

bool QueryHasTextPredicate(const SearchParameters &parameters) {
  return parameters.filter_parse_results.query_operations &
         QueryOperations::kContainsText;
}

// Increment query operation metrics based on query operations flags.
// File-internal helper function.
void IncrementQueryOperationMetrics(QueryOperations query_operations) {
  // High-level query type metrics
  if (query_operations & QueryOperations::kContainsText) {
    ++Metrics::GetStats().query_text_requests_cnt;
  }
  if (query_operations & QueryOperations::kContainsNumeric) {
    query_numeric_count.Increment();
  }
  if (query_operations & QueryOperations::kContainsTag) {
    query_tag_count.Increment();
  }
  // Text operation type metrics
  if (query_operations & QueryOperations::kContainsTextTerm) {
    query_text_term_count.Increment();
  }
  if (query_operations & QueryOperations::kContainsTextPrefix) {
    query_text_prefix_count.Increment();
  }
  if (query_operations & QueryOperations::kContainsTextSuffix) {
    query_text_suffix_count.Increment();
  }
  if (query_operations & QueryOperations::kContainsTextFuzzy) {
    query_text_fuzzy_count.Increment();
  }
  if (query_operations & QueryOperations::kContainsProximity) {
    query_text_proximity_count.Increment();
  }
}

absl::StatusOr<absl::string_view> SubstituteParam(
    query::SearchParameters &parameters, absl::string_view source) {
  if (source.empty() || source[0] != '$') {
    return source;
  } else {
    source.remove_prefix(1);
    auto itr = parameters.parse_vars.params.find(source);
    if (itr == parameters.parse_vars.params.end()) {
      return absl::NotFoundError(
          absl::StrCat("Parameter ", source, " not found."));
    } else {
      itr->second.first++;
      return itr->second.second;
    }
  }
}

absl::Status ParseKnnInner(query::SearchParameters &parameters,
                           std::string_view filter) {
  absl::InlinedVector<absl::string_view, 8> params =
      absl::StrSplit(filter, ' ', absl::SkipEmpty());
  if (params.empty()) {
    return absl::InvalidArgumentError("Missing parameters");
  }
  // TODO - need some investment to consolidate this with the common parsing
  // functionality
  if (!absl::EqualsIgnoreCase(params[0], "KNN")) {
    return absl::InvalidArgumentError(
        absl::StrCat("`", params[0], "`. Expecting `KNN`"));
  }
  if (params.size() == 1) {
    return absl::InvalidArgumentError("KNN argument is missing");
  }
  parameters.parse_vars.k_string = params[1];
  if (params.size() == 2) {
    return absl::InvalidArgumentError("Vector field argument is missing");
  }
  if (params[2].data()[0] != '@' || params[2].size() == 1) {
    return absl::InvalidArgumentError(
        absl::StrCat("Unexpected argument `", params[2],
                     "`. Expecting a vector field name, starting with '@'"));
  }
  parameters.attribute_alias =
      absl::string_view(params[2].data() + 1, params[2].size() - 1);
  if (params.size() == 3) {
    return absl::InvalidArgumentError("Blob attribute argument is missing");
  }
  parameters.parse_vars.query_vector_string = params[3];

  size_t i = 4;
  while (i < params.size()) {
    if (absl::EqualsIgnoreCase(params[i], "EF_RUNTIME")) {
      i++;
      if (i == params.size()) {
        return absl::InvalidArgumentError("EF_RUNTIME argument is missing");
      }
      parameters.parse_vars.ef_string = params[i++];
    } else if (absl::EqualsIgnoreCase(params[i], kAsParam)) {
      i++;
      if (i == params.size()) {
        return absl::InvalidArgumentError("AS argument is missing");
      }
      parameters.parse_vars.score_as_string = params[i++];
    } else {
      return absl::InvalidArgumentError(
          absl::StrCat("Unexpected argument `", params[i], "`"));
    }
  }
  return absl::OkStatus();
}

absl::StatusOr<size_t> FindOpenSquareBracket(absl::string_view input) {
  for (size_t position = 0; position < input.size(); ++position) {
    if (input[position] == '[') {
      return position;
    }
    if (!std::isspace(input[position])) {
      return absl::InvalidArgumentError(
          absl::StrCat("Expecting '[' got '", input.substr(position, 1), "'"));
    }
  }
  return absl::InvalidArgumentError("Missing opening bracket");
}

absl::StatusOr<size_t> FindCloseSquareBracket(absl::string_view input) {
  for (auto position = input.size(); position > 0; --position) {
    if (input[position - 1] == ']') {
      return position - 1;
    }
    if (!std::isspace(input[position - 1])) {
      return absl::InvalidArgumentError(absl::StrCat(
          "Expecting ']' got '", input.substr(position - 1, 1), "'"));
    }
  }
  if (input[0] == ']') {
    return 0;
  }
  return absl::InvalidArgumentError("Missing closing bracket");
}

absl::StatusOr<FilterParseResults> ParsePreFilter(
    const IndexSchema &index_schema, absl::string_view pre_filter,
    const query::SearchParameters &search_params) {
  TextParsingOptions options{.verbatim = search_params.verbatim,
                             .inorder = search_params.inorder,
                             .slop = search_params.slop};
  FilterParser parser(index_schema, pre_filter, options);
  return parser.Parse();
}

absl::Status ParseKNN(query::SearchParameters &parameters,
                      absl::string_view filter_str) {
  if (filter_str.empty()) {
    return absl::InvalidArgumentError("Vector query clause is missing");
  }
  VMSDK_ASSIGN_OR_RETURN(auto close_position,
                         FindCloseSquareBracket(filter_str));
  size_t position = 0;
  VMSDK_ASSIGN_OR_RETURN(
      auto open_position,
      FindOpenSquareBracket(absl::string_view(filter_str.data() + position,
                                              close_position - position)));
  position += open_position;
  return ParseKnnInner(parameters,
                       absl::string_view(filter_str.data() + position + 1,
                                         close_position - position - 1));
}

//
// We don't have values for the $ substitution yet. so we break the parsing into
// two pieces
//
absl::Status query::SearchParameters::PreParseQueryString() {
  // Validate the query string's length.
  if (parse_vars.query_string.length() > options::GetQueryStringBytes()) {
    return absl::InvalidArgumentError(
        absl::StrCat("Query string is too long, max length is ",
                     options::GetQueryStringBytes(), " bytes."));
  }
  auto filter_expression = absl::string_view(parse_vars.query_string);
  VMSDK_LOG(DEBUG, nullptr)
      << "Query: '" << vmsdk::config::RedactIfNeeded(parse_vars.query_string)
      << "'";
  auto pos = filter_expression.find(kVectorFilterDelimiter);
  absl::string_view pre_filter;
  absl::string_view vector_filter;
  // If the delimiter is not found (ie - non vector query), treat the whole
  // string as pre-filter.
  if (pos == absl::string_view::npos) {
    pre_filter = absl::StripAsciiWhitespace(filter_expression);
  } else {
    pre_filter = absl::StripAsciiWhitespace(filter_expression.substr(0, pos));
    vector_filter = absl::StripAsciiWhitespace(
        filter_expression.substr(pos + kVectorFilterDelimiter.size()));
  }
  // If INORDER OR SLOP, but the index schema does not support offsets, we
  // reject the query.
  if ((inorder || slop.has_value()) && !index_schema->HasTextOffsets()) {
    return absl::InvalidArgumentError("Index does not support offsets");
  }
  VMSDK_ASSIGN_OR_RETURN(
      filter_parse_results, ParsePreFilter(*index_schema, pre_filter, *this),
      _.SetPrepend() << "Invalid filter expression: `" << pre_filter << "`. ");
  if (!filter_parse_results.root_predicate && vector_filter.empty()) {
    // Return an error if no valid pre-filter and no vector filter is provided.
    return absl::InvalidArgumentError("Invalid query string syntax");
  }
  // Optionally parse the vector filter if it exists.
  if (!vector_filter.empty()) {
    if (filter_parse_results.root_predicate) {
      ++Metrics::GetStats().query_hybrid_requests_cnt;
    } else {
      // Pure vector query
      ++Metrics::GetStats().query_vector_requests_cnt;
    }
    VMSDK_RETURN_IF_ERROR(ParseKNN(*this, vector_filter)).SetPrepend()
        << "Error parsing vector similarity parameters: `" << vector_filter
        << "`. ";
    // Validate the index exists and is a vector index.
    VMSDK_ASSIGN_OR_RETURN(auto index, index_schema->GetIndex(attribute_alias));
    if (index->GetIndexerType() != indexes::IndexerType::kHNSW &&
        index->GetIndexerType() != indexes::IndexerType::kFlat) {
      return absl::InvalidArgumentError(absl::StrCat(
          "Index field `", attribute_alias, "` is not a Vector index "));
    }
    if (parse_vars.score_as_string.empty()) {
      VMSDK_ASSIGN_OR_RETURN(
          score_as, index_schema->DefaultReplyScoreAs(attribute_alias));
    } else {
      score_as = vmsdk::MakeUniqueValkeyString(parse_vars.score_as_string);
    }
  }

  // Pure non-vector query (no vector filter)
  if (vector_filter.empty() && filter_parse_results.root_predicate) {
    ++Metrics::GetStats().query_nonvector_requests_cnt;
  }
  // Increment operation-type metrics
  IncrementQueryOperationMetrics(filter_parse_results.query_operations);
  return absl::OkStatus();
}

absl::Status PostParseVectorParameters(query::SearchParameters &parameters) {
  VMSDK_ASSIGN_OR_RETURN(
      auto k_string,
      SubstituteParam(parameters, parameters.parse_vars.k_string));
  VMSDK_ASSIGN_OR_RETURN(parameters.k, vmsdk::To<unsigned>(k_string));

  VMSDK_ASSIGN_OR_RETURN(
      parameters.query,
      SubstituteParam(parameters, parameters.parse_vars.query_vector_string));

  if (!parameters.parse_vars.ef_string.empty()) {
    VMSDK_ASSIGN_OR_RETURN(
        auto ef_string,
        SubstituteParam(parameters, parameters.parse_vars.ef_string));
    VMSDK_ASSIGN_OR_RETURN(parameters.ef, vmsdk::To<unsigned>(ef_string));
  }

  if (!parameters.parse_vars.score_as_string.empty()) {
    VMSDK_ASSIGN_OR_RETURN(
        parameters.parse_vars.score_as_string,
        SubstituteParam(parameters, parameters.parse_vars.score_as_string));
  }
  return absl::OkStatus();
}

absl::Status query::SearchParameters::PostParseQueryString() {
  if (IsVectorQuery()) {
    VMSDK_RETURN_IF_ERROR(PostParseVectorParameters(*this)).SetPrepend()
        << "Error parsing vector similarity parameters: ";
  }

  return absl::OkStatus();
}

ContentProcessing SearchParameters::GetContentProcessing() const {
  if (no_content) {
    return kNoContent;
  }
  // Currently, ContentAvailable isn't detected. Future use case.
  if (query::QueryHasTextPredicate(*this)) {
    return kContentionCheckRequired;
  }
  return kContentRequired;
}

}  // namespace valkey_search::query
