/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */
package org.apache.lucene.search;

import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
import static org.apache.lucene.search.AbstractVectorSimilarityQuery.DECAY_MAX_APPROXIMATION;
import static org.apache.lucene.search.AbstractVectorSimilarityQuery.DECAY_MAX_QUALITY;
import static org.apache.lucene.search.AbstractVectorSimilarityQuery.DEFAULT_DECAY;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.IntStream;
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.IntField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.QueryTimeout;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.search.knn.KnnSearchStrategy;
import org.apache.lucene.store.Directory;
import org.apache.lucene.tests.index.RandomIndexWriter;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.hnsw.HnswUtil;

@LuceneTestCase.SuppressCodecs("SimpleText")
abstract class BaseVectorSimilarityQueryTestCase<
        V, F extends Field, Q extends AbstractVectorSimilarityQuery>
    extends LuceneTestCase {
  String vectorField, idField;
  VectorSimilarityFunction function;
  int numDocs, dim;

  abstract V getRandomVector(int dim);

  abstract float compare(V vector1, V vector2);

  abstract boolean checkEquals(V vector1, V vector2);

  abstract F getVectorField(String name, V vector, VectorSimilarityFunction function);

  abstract Q getVectorQuery(
      String field, V vector, float resultSimilarity, float decay, Query filter);

  abstract Q getVectorQuery(
      String field,
      V vector,
      float resultSimilarity,
      float decay,
      Query filter,
      KnnSearchStrategy searchStrategy);

  abstract Q getThrowingVectorQuery(
      String field, V vector, float resultSimilarity, float decay, Query filter);

  public void testEquals() {
    String field1 = "f1", field2 = "f2";

    V vector1 = getRandomVector(dim);
    V vector2;
    do {
      vector2 = getRandomVector(dim);
    } while (checkEquals(vector1, vector2));

    float resultSimilarity1 = 0.4f, resultSimilarity2 = 0.5f;
    float decay1 = 0.5f, decay2 = 0.6f;

    Query filter1 = new TermQuery(new Term("t1", "v1"));
    Query filter2 = new TermQuery(new Term("t2", "v2"));

    Query query = getVectorQuery(field1, vector1, resultSimilarity1, decay1, filter1);

    // Everything is equal
    assertEquals(query, getVectorQuery(field1, vector1, resultSimilarity1, decay1, filter1));

    // Null check
    assertNotEquals(null, query);

    // Different field
    assertNotEquals(query, getVectorQuery(field2, vector1, resultSimilarity1, decay1, filter1));

    // Different vector
    assertNotEquals(query, getVectorQuery(field1, vector2, resultSimilarity1, decay1, filter1));

    // Different decay
    assertNotEquals(query, getVectorQuery(field1, vector1, resultSimilarity1, decay2, filter1));

    // Different resultSimilarity
    assertNotEquals(query, getVectorQuery(field1, vector1, resultSimilarity2, decay1, filter1));

    // Different filter
    assertNotEquals(query, getVectorQuery(field1, vector1, resultSimilarity1, decay1, filter2));
  }

  public void testEqualsWithSearchStrategy() {
    String field = "f";
    V vector = getRandomVector(dim);
    float resultSimilarity = 0.4f;
    float decay = 0.5f;
    Query filter = new TermQuery(new Term("t", "v"));

    KnnSearchStrategy strategyA = new KnnSearchStrategy.Hnsw(10);
    KnnSearchStrategy strategyB = new KnnSearchStrategy.Hnsw(20);

    Q queryA = getVectorQuery(field, vector, resultSimilarity, decay, filter, strategyA);
    Q queryB = getVectorQuery(field, vector, resultSimilarity, decay, filter, strategyB);
    Q queryADup = getVectorQuery(field, vector, resultSimilarity, decay, filter, strategyA);

    // Queries with the same non-default strategy are equal.
    assertEquals(queryA, queryADup);
    assertEquals(queryA.hashCode(), queryADup.hashCode());

    // Queries that differ only in strategy are not equal.
    assertNotEquals(queryA, queryB);

    // Query with default strategy is not equal to one with non-default strategy.
    Q queryDefault = getVectorQuery(field, vector, resultSimilarity, decay, filter);
    assertNotEquals(queryDefault, queryA);
  }

  public void testGetSearchStrategy() {
    String field = "f";
    V vector = getRandomVector(dim);
    KnnSearchStrategy strategy = new KnnSearchStrategy.Hnsw(42);

    // Default strategy is exposed via getter.
    AbstractVectorSimilarityQuery defaultQuery = getVectorQuery(field, vector, 0.4f, 0.5f, null);
    assertEquals(AbstractVectorSimilarityQuery.DEFAULT_STRATEGY, defaultQuery.getSearchStrategy());

    // Custom strategy is propagated through to the getter.
    AbstractVectorSimilarityQuery customQuery =
        getVectorQuery(field, vector, 0.4f, 0.5f, null, strategy);
    assertSame(strategy, customQuery.getSearchStrategy());
  }

  public void testNullSearchStrategyDefaultsToDefault() {
    String field = "f";
    V vector = getRandomVector(dim);
    Q query = getVectorQuery(field, vector, 0.4f, 0.5f, null, /* searchStrategy= */ null);
    assertSame(AbstractVectorSimilarityQuery.DEFAULT_STRATEGY, query.getSearchStrategy());
  }

  public void testSearchStrategyReachesCollector() throws IOException {
    KnnSearchStrategy strategy = new KnnSearchStrategy.Hnsw(7);
    Q query =
        getVectorQuery(
            "field",
            getRandomVector(dim),
            /* resultSimilarity= */ 0.5f,
            /* decay= */ 0.5f,
            /* filter= */ null,
            strategy);

    KnnCollector collector =
        query
            .getKnnCollectorManager()
            .newCollector(Integer.MAX_VALUE, query.getSearchStrategy(), null);

    assertSame(strategy, collector.getSearchStrategy());
  }

  public void testIllegalParams() {
    expectThrows(
        IllegalArgumentException.class,
        () ->
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                Float.NaN, // illegal resultSimilarity
                DEFAULT_DECAY,
                null));

    expectThrows(
        IllegalArgumentException.class,
        () ->
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                0f,
                Float.NaN, // illegal decay
                null));

    expectThrows(
        IllegalArgumentException.class,
        () ->
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                0f,
                Math.nextDown(DECAY_MAX_APPROXIMATION), // illegal decay
                null));

    expectThrows(
        IllegalArgumentException.class,
        () ->
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                0f,
                Math.nextUp(DECAY_MAX_QUALITY), // illegal decay
                null));
  }

  public void testEmptyIndex() throws IOException {
    // Do not index any vectors
    numDocs = 0;

    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      Query query =
          getVectorQuery(
              vectorField, getRandomVector(dim), Float.NEGATIVE_INFINITY, DECAY_MAX_QUALITY, null);

      // Check that no vectors are found
      assertEquals(0, searcher.count(query));
    }
  }

  public void testExtremes() throws IOException {
    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);
      assumeTrue("graph is disconnected", HnswUtil.graphIsRooted(reader, vectorField));

      // All vectors are above -Infinity
      Query query1 =
          getVectorQuery(
              vectorField, getRandomVector(dim), Float.NEGATIVE_INFINITY, DECAY_MAX_QUALITY, null);

      // Check that all vectors are found
      assertEquals(numDocs, searcher.count(query1));

      // No vectors are above +Infinity
      Query query2 =
          getVectorQuery(
              vectorField, getRandomVector(dim), Float.POSITIVE_INFINITY, DECAY_MAX_QUALITY, null);

      // Check that no vectors are found
      assertEquals(0, searcher.count(query2));
    }
  }

  public void testRandomFilter() throws IOException {
    // Filter a sub-range from 0 to numDocs
    int startIndex = random().nextInt(numDocs);
    int endIndex = random().nextInt(startIndex, numDocs);
    Query filter = IntField.newRangeQuery(idField, startIndex, endIndex);

    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      assumeTrue("graph is disconnected", HnswUtil.graphIsRooted(reader, vectorField));
      IndexSearcher searcher = newSearcher(reader);

      Query query =
          getVectorQuery(
              vectorField,
              getRandomVector(dim),
              Float.NEGATIVE_INFINITY,
              DECAY_MAX_QUALITY,
              filter);

      ScoreDoc[] scoreDocs = searcher.search(query, numDocs).scoreDocs;
      for (ScoreDoc scoreDoc : scoreDocs) {
        int id = getId(searcher, scoreDoc.doc);

        // Check that returned document is in selected range
        assertTrue(id >= startIndex && id <= endIndex);
      }
      // Check that all filtered vectors are found
      assertEquals(endIndex - startIndex + 1, scoreDocs.length);
    }
  }

  public void testFilterWithNoMatches() throws IOException {
    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      // Non-existent field
      Query filter1 = new TermQuery(new Term("random_field", "random_value"));
      Query query1 =
          getVectorQuery(
              vectorField,
              getRandomVector(dim),
              Float.NEGATIVE_INFINITY,
              DECAY_MAX_QUALITY,
              filter1);

      // Check that no vectors are found
      assertEquals(0, searcher.count(query1));

      // Field exists, but value of -1 is not indexed
      Query filter2 = IntField.newExactQuery(idField, -1);
      Query query2 =
          getVectorQuery(
              vectorField,
              getRandomVector(dim),
              Float.NEGATIVE_INFINITY,
              DECAY_MAX_QUALITY,
              filter2);

      // Check that no vectors are found
      assertEquals(0, searcher.count(query2));
    }
  }

  public void testDimensionMismatch() throws IOException {
    // Different dimension
    int newDim = atLeast(dim + 1);

    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      Query query =
          getVectorQuery(
              vectorField,
              getRandomVector(newDim),
              Float.NEGATIVE_INFINITY,
              DECAY_MAX_QUALITY,
              null);

      // Check that an exception for differing dimensions is thrown
      IllegalArgumentException e =
          expectThrows(IllegalArgumentException.class, () -> searcher.count(query));
      assertEquals(
          String.format(
              Locale.ROOT,
              "vector query dimension: %d differs from field dimension: %d",
              newDim,
              dim),
          e.getMessage());
    }
  }

  public void testNonVectorsField() throws IOException {
    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      // Non-existent field
      Query query1 =
          getVectorQuery(
              "random_field",
              getRandomVector(dim),
              Float.NEGATIVE_INFINITY,
              DECAY_MAX_QUALITY,
              null);
      assertEquals(0, searcher.count(query1));

      // Indexed as int field
      Query query2 =
          getVectorQuery(
              idField, getRandomVector(dim), Float.NEGATIVE_INFINITY, DECAY_MAX_QUALITY, null);
      assertEquals(0, searcher.count(query2));
    }
  }

  public void testSomeDeletes() throws IOException {
    // Delete a sub-range from 0 to numDocs
    int startIndex = random().nextInt(numDocs);
    int endIndex = random().nextInt(startIndex, numDocs);
    Query delete = IntField.newRangeQuery(idField, startIndex, endIndex);

    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexWriter w = new IndexWriter(indexStore, newIndexWriterConfig())) {

      w.deleteDocuments(delete);
      w.commit();

      try (IndexReader reader = DirectoryReader.open(indexStore)) {
        assumeTrue("graph is disconnected", HnswUtil.graphIsRooted(reader, vectorField));
        IndexSearcher searcher = newSearcher(reader);

        Query query =
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                Float.NEGATIVE_INFINITY,
                DECAY_MAX_QUALITY,
                null);

        ScoreDoc[] scoreDocs = searcher.search(query, numDocs).scoreDocs;
        for (ScoreDoc scoreDoc : scoreDocs) {
          int id = getId(searcher, scoreDoc.doc);

          // Check that returned document is not deleted
          assertFalse(id >= startIndex && id <= endIndex);
        }
        // Check that all live docs are returned
        assertEquals(numDocs - endIndex + startIndex - 1, scoreDocs.length);
      }
    }
  }

  public void testAllDeletes() throws IOException {
    try (Directory dir = getIndexStore(getRandomVectors(numDocs, dim));
        IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) {
      // Delete all documents
      w.deleteDocuments(MatchAllDocsQuery.INSTANCE);
      w.commit();

      try (IndexReader reader = DirectoryReader.open(dir)) {
        IndexSearcher searcher = newSearcher(reader);

        Query query =
            getVectorQuery(
                vectorField,
                getRandomVector(dim),
                Float.NEGATIVE_INFINITY,
                DECAY_MAX_QUALITY,
                null);

        // Check that no vectors are found
        assertEquals(0, searcher.count(query));
      }
    }
  }

  public void testBoostQuery() throws IOException {
    // Define the boost and allowed delta
    float boost = random().nextFloat(5, 10);
    float delta = 1e-3f;

    try (Directory indexStore = getIndexStore(getRandomVectors(numDocs, dim));
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      Query query1 =
          getVectorQuery(
              vectorField, getRandomVector(dim), Float.NEGATIVE_INFINITY, DECAY_MAX_QUALITY, null);
      ScoreDoc[] scoreDocs1 = searcher.search(query1, numDocs).scoreDocs;

      Query query2 = new BoostQuery(query1, boost);
      ScoreDoc[] scoreDocs2 = searcher.search(query2, numDocs).scoreDocs;

      // Check that original scores and boosted scores are equal considering the delta to account
      // for floating point precision limitations. Don't take the exact result order into
      // consideration as for small original scores with tiny differences
      // the boosted scores can be become the same, which might affect the result order.
      assertEquals(scoreDocs1.length, scoreDocs2.length);
      for (int i = 0; i < scoreDocs1.length; i++) {
        int idx = i;
        Optional<ScoreDoc> boostedDoc =
            Arrays.stream(scoreDocs2).filter(d -> d.doc == scoreDocs1[idx].doc).findFirst();

        assertTrue(boostedDoc.isPresent());
        assertEquals(boost * scoreDocs1[i].score, boostedDoc.get().score, delta);
      }
    }
  }

  public void testVectorsAboveSimilarity() throws IOException {
    // Pick number of docs to accept
    int numAccepted = random().nextInt(numDocs / 3, numDocs / 2);
    float delta = 1e-3f;

    V[] vectors = getRandomVectors(numDocs, dim);
    V queryVector = getRandomVector(dim);

    // Find score above which we get (at least) numAccepted vectors
    float resultSimilarity = getSimilarity(vectors, queryVector, numAccepted);

    // Cache scores of vectors
    Map<Integer, Float> scores = new HashMap<>();
    for (int i = 0; i < numDocs; i++) {
      float score = compare(queryVector, vectors[i]);
      if (score >= resultSimilarity) {
        scores.put(i, score);
      }
    }

    try (Directory indexStore = getIndexStore(vectors);
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      Query query =
          getVectorQuery(vectorField, queryVector, resultSimilarity, DECAY_MAX_QUALITY, null);

      ScoreDoc[] scoreDocs = searcher.search(query, numDocs).scoreDocs;
      for (ScoreDoc scoreDoc : scoreDocs) {
        int id = getId(searcher, scoreDoc.doc);

        // Check that the collected result is above accepted similarity
        assertTrue(scores.containsKey(id));

        // Check that the score is correct
        assertEquals(scores.get(id), scoreDoc.score, delta);
      }

      // Check that all results are collected
      assertEquals(scores.size(), scoreDocs.length);
    }
  }

  public void testFallbackToExact() throws IOException {
    // Restrictive filter, along with similarity to visit a large number of nodes
    int numFiltered = numDocs / 5;
    int targetVisited = numDocs;

    V[] vectors = getRandomVectors(numDocs, dim);
    V queryVector = getRandomVector(dim);

    float resultSimilarity = getSimilarity(vectors, queryVector, targetVisited);
    Query filter = IntField.newSetQuery(idField, getFiltered(numFiltered));

    try (Directory indexStore = getIndexStore(vectors);
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      Query query =
          getThrowingVectorQuery(
              vectorField, queryVector, resultSimilarity, DECAY_MAX_APPROXIMATION, filter);

      // Falls back to exact search, even with DECAY_MAX_APPROXIMATION
      expectThrows(UnsupportedOperationException.class, () -> searcher.count(query));

      Query exactQuery =
          getThrowingVectorQuery(
              vectorField, queryVector, Float.NEGATIVE_INFINITY, DECAY_MAX_QUALITY, null);

      // Exact search should be used directly with DECAY_MAX_QUALITY
      expectThrows(UnsupportedOperationException.class, () -> searcher.count(exactQuery));
    }
  }

  @Monster("indexes and searches a large number of vectors")
  public void testApproximate() throws IOException {
    numDocs = 1000;

    // Non-restrictive filter, along with similarity to visit a small number of nodes
    int numFiltered = numDocs - 1;
    int targetVisited = random().nextInt(1, numFiltered / 10);

    V[] vectors = getRandomVectors(numDocs, dim);
    V queryVector = getRandomVector(dim);

    float resultSimilarity = getSimilarity(vectors, queryVector, targetVisited);
    Query filter = IntField.newSetQuery(idField, getFiltered(numFiltered));

    try (Directory indexStore = getIndexStore(vectors);
        IndexWriter w =
            new IndexWriter(
                indexStore,
                newIndexWriterConfig()
                    .setCodec(
                        TestUtil.alwaysKnnVectorsFormat(
                            new Lucene99HnswVectorsFormat(
                                DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, 0))))) {
      // Force merge because smaller segments have few filtered docs and often fall back to exact
      // search, making this test flaky
      w.forceMerge(1);
      w.commit();

      try (IndexReader reader = DirectoryReader.open(indexStore)) {
        IndexSearcher searcher = newSearcher(reader);

        // Use DECAY_MAX_APPROXIMATION to avoid fallback to exact search
        Query query =
            getThrowingVectorQuery(
                vectorField, queryVector, resultSimilarity, DECAY_MAX_APPROXIMATION, filter);

        // Does not fall back to exact search
        assertTrue(searcher.count(query) <= numFiltered);
      }
    }
  }

  /** Test that the query times out correctly. */
  public void testTimeout() throws IOException {
    V[] vectors = getRandomVectors(numDocs, dim);
    V queryVector = getRandomVector(dim);

    try (Directory indexStore = getIndexStore(vectors);
        IndexReader reader = DirectoryReader.open(indexStore)) {
      IndexSearcher searcher = newSearcher(reader);

      // This query is cacheable, explicitly prevent it
      searcher.setQueryCache(null);

      // Use Math.nextDown(DECAY_MAX_QUALITY) to ensure approximate graph search is used,
      // since DECAY_MAX_QUALITY now triggers exact search which bypasses graph traversal
      float highQualityNoExact = Math.nextDown(DECAY_MAX_QUALITY);

      Query query =
          new CountingQuery(
              getVectorQuery(
                  vectorField, queryVector, Float.NEGATIVE_INFINITY, highQualityNoExact, null));
      assertEquals(numDocs, searcher.count(query)); // Expect some results without timeout

      searcher.setTimeout(() -> true); // Immediately timeout
      assertEquals(0, searcher.count(query)); // Expect no results with the timeout

      searcher.setTimeout(new CountingQueryTimeout(numDocs - 1)); // Do not score all docs
      int count = searcher.count(query);
      assertTrue(
          "0 < count=" + count + " < numDocs=" + numDocs,
          count > 0 && count < numDocs); // Expect partial results

      // Test timeout with filter
      int numFiltered = random().nextInt(numDocs / 2, numDocs);
      Query filter = IntField.newSetQuery(idField, getFiltered(numFiltered));
      Query filteredQuery =
          new CountingQuery(
              getVectorQuery(
                  vectorField, queryVector, Float.NEGATIVE_INFINITY, highQualityNoExact, filter));

      searcher.setTimeout(() -> false); // Set a timeout which is never met
      assertEquals(numFiltered, searcher.count(filteredQuery));

      searcher.setTimeout(
          new CountingQueryTimeout(numFiltered - 1)); // Timeout before scoring all filtered docs
      int filteredCount = searcher.count(filteredQuery);
      assertTrue(
          "0 < filteredCount=" + filteredCount + " < numFiltered=" + numFiltered,
          filteredCount > 0 && filteredCount < numFiltered); // Expect partial results
    }
  }

  private float getSimilarity(V[] vectors, V queryVector, int targetVisited) {
    assertTrue(targetVisited >= 0 && targetVisited <= numDocs);
    if (targetVisited == 0) {
      return Float.POSITIVE_INFINITY;
    }

    float[] scores = new float[numDocs];
    for (int i = 0; i < numDocs; i++) {
      scores[i] = compare(queryVector, vectors[i]);
    }
    Arrays.sort(scores);

    return scores[numDocs - targetVisited];
  }

  private int[] getFiltered(int numFiltered) {
    Set<Integer> accepted = new HashSet<>();
    for (int i = 0; i < numFiltered; ) {
      int index = random().nextInt(numDocs);
      if (!accepted.contains(index)) {
        accepted.add(index);
        i++;
      }
    }
    return accepted.stream().mapToInt(Integer::intValue).toArray();
  }

  private int getId(IndexSearcher searcher, int doc) throws IOException {
    return Objects.requireNonNull(searcher.storedFields().document(doc).getField(idField))
        .numericValue()
        .intValue();
  }

  @SuppressWarnings("unchecked")
  V[] getRandomVectors(int numDocs, int dim) {
    return (V[]) IntStream.range(0, numDocs).mapToObj(i -> getRandomVector(dim)).toArray();
  }

  @SafeVarargs
  final Directory getIndexStore(V... vectors) throws IOException {
    Directory dir = newDirectory();
    try (RandomIndexWriter writer =
        new RandomIndexWriter(
            random(),
            dir,
            newIndexWriterConfig()
                .setCodec(
                    TestUtil.alwaysKnnVectorsFormat(
                        new Lucene99HnswVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, 0))))) {
      for (int i = 0; i < vectors.length; ++i) {
        Document doc = new Document();
        doc.add(getVectorField(vectorField, vectors[i], function));
        doc.add(new IntField(idField, i, Field.Store.YES));
        writer.addDocument(doc);
      }
    }
    return dir;
  }

  private static class CountingQueryTimeout implements QueryTimeout {
    private int remaining;

    public CountingQueryTimeout(int count) {
      remaining = count;
    }

    @Override
    public boolean shouldExit() {
      if (remaining > 0) {
        remaining--;
        return false;
      }
      return true;
    }
  }

  /**
   * A {@link Query} that emulates {@link Weight#count(LeafReaderContext)} by counting number of
   * docs of underlying {@link Scorer#iterator()}. TODO: This is a workaround to count partial
   * results of {@link #delegate} because {@link TimeLimitingBulkScorer} immediately discards
   * results after timeout.
   */
  private static class CountingQuery extends Query {
    private final Query delegate;

    private CountingQuery(Query delegate) {
      this.delegate = delegate;
    }

    @Override
    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost)
        throws IOException {
      return new Weight(this) {
        final Weight delegateWeight = delegate.createWeight(searcher, scoreMode, boost);

        @Override
        public Explanation explain(LeafReaderContext context, int doc) throws IOException {
          return delegateWeight.explain(context, doc);
        }

        @Override
        public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
          return delegateWeight.scorerSupplier(context);
        }

        @Override
        public int count(LeafReaderContext context) throws IOException {
          Scorer scorer = scorer(context);
          if (scorer == null) {
            return 0;
          }

          int count = 0;
          DocIdSetIterator iterator = scorer.iterator();
          while (iterator.nextDoc() != NO_MORE_DOCS) {
            count++;
          }
          return count;
        }

        @Override
        public boolean isCacheable(LeafReaderContext ctx) {
          return delegateWeight.isCacheable(ctx);
        }
      };
    }

    @Override
    public String toString(String field) {
      return String.format(
          Locale.ROOT, "%s[%s]", getClass().getSimpleName(), delegate.toString(field));
    }

    @Override
    public void visit(QueryVisitor visitor) {
      visitor.visitLeaf(this);
    }

    @Override
    public boolean equals(Object obj) {
      return sameClassAs(obj) && delegate.equals(((CountingQuery) obj).delegate);
    }

    @Override
    public int hashCode() {
      return delegate.hashCode();
    }
  }
}
