/*
 * 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.codecs.lucene99;

import static org.apache.lucene.codecs.KnnVectorsWriter.MergedVectorValues.hasVectorValues;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DIRECT_MONOTONIC_BLOCK_SHIFT;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.HNSW_GRAPH_THRESHOLD;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.VERSION_CURRENT;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.VERSION_GROUPVARINT;
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
import static org.apache.lucene.util.hnsw.HnswGraphSearcher.expectedVisitedNodes;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.KnnFieldVectorsWriter;
import org.apache.lucene.codecs.KnnVectorsWriter;
import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter;
import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
import org.apache.lucene.index.DocsWithFieldSet;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.IndexFileNames;
import org.apache.lucene.index.KnnVectorValues;
import org.apache.lucene.index.MergeState;
import org.apache.lucene.index.SegmentReadState;
import org.apache.lucene.index.SegmentWriteState;
import org.apache.lucene.index.Sorter;
import org.apache.lucene.index.VectorEncoding;
import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.search.TaskExecutor;
import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.util.IORunnable;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.InfoStream;
import org.apache.lucene.util.RamUsageEstimator;
import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier;
import org.apache.lucene.util.hnsw.ConcurrentHnswMerger;
import org.apache.lucene.util.hnsw.HnswGraph;
import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator;
import org.apache.lucene.util.hnsw.HnswGraphBuilder;
import org.apache.lucene.util.hnsw.HnswGraphMerger;
import org.apache.lucene.util.hnsw.IncrementalHnswGraphMerger;
import org.apache.lucene.util.hnsw.NeighborArray;
import org.apache.lucene.util.hnsw.OnHeapHnswGraph;
import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
import org.apache.lucene.util.packed.DirectMonotonicWriter;
import org.apache.lucene.util.quantization.QuantizedVectorsReader;

/**
 * Writes vector values and knn graphs to index segments.
 *
 * @lucene.experimental
 */
public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {

  private static final long SHALLOW_RAM_BYTES_USED =
      RamUsageEstimator.shallowSizeOfInstance(Lucene99HnswVectorsWriter.class);
  private final SegmentWriteState segmentWriteState;
  private final IndexOutput meta, vectorIndex;
  private final int M;
  private final int beamWidth;
  private final FlatVectorsWriter flatVectorWriter;
  private final FlatVectorsFormat flatVectorsFormat;
  private FlatVectorsReader flatVectorsReader;
  private boolean flatWriterClosed = false;
  private final int numMergeWorkers;
  private final TaskExecutor mergeExec;
  private final int tinySegmentsThreshold;
  private final int version;

  private final List<FieldWriter<?>> fields = new ArrayList<>();
  private boolean finished;

  public Lucene99HnswVectorsWriter(
      SegmentWriteState state,
      int M,
      int beamWidth,
      FlatVectorsFormat flatVectorsFormat,
      FlatVectorsWriter flatVectorWriter,
      int numMergeWorkers,
      TaskExecutor mergeExec)
      throws IOException {
    this(
        state,
        M,
        beamWidth,
        flatVectorsFormat,
        flatVectorWriter,
        numMergeWorkers,
        mergeExec,
        HNSW_GRAPH_THRESHOLD,
        VERSION_CURRENT);
  }

  public Lucene99HnswVectorsWriter(
      SegmentWriteState state,
      int M,
      int beamWidth,
      FlatVectorsFormat flatVectorsFormat,
      FlatVectorsWriter flatVectorWriter,
      int numMergeWorkers,
      TaskExecutor mergeExec,
      int tinySegmentsThreshold)
      throws IOException {
    this(
        state,
        M,
        beamWidth,
        flatVectorsFormat,
        flatVectorWriter,
        numMergeWorkers,
        mergeExec,
        tinySegmentsThreshold,
        VERSION_CURRENT);
  }

  Lucene99HnswVectorsWriter(
      SegmentWriteState state,
      int M,
      int beamWidth,
      FlatVectorsFormat flatVectorsFormat,
      FlatVectorsWriter flatVectorWriter,
      int numMergeWorkers,
      TaskExecutor mergeExec,
      int tinySegmentsThreshold,
      int version)
      throws IOException {
    this.M = M;
    this.flatVectorWriter = flatVectorWriter;
    this.flatVectorsFormat = flatVectorsFormat;
    this.beamWidth = beamWidth;
    this.numMergeWorkers = numMergeWorkers;
    this.mergeExec = mergeExec;
    this.tinySegmentsThreshold = tinySegmentsThreshold;
    this.version = version;
    segmentWriteState = state;

    String metaFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name, state.segmentSuffix, Lucene99HnswVectorsFormat.META_EXTENSION);

    String indexDataFileName =
        IndexFileNames.segmentFileName(
            state.segmentInfo.name,
            state.segmentSuffix,
            Lucene99HnswVectorsFormat.VECTOR_INDEX_EXTENSION);

    boolean success = false;
    try {
      meta = state.directory.createOutput(metaFileName, state.context);
      vectorIndex = state.directory.createOutput(indexDataFileName, state.context);

      CodecUtil.writeIndexHeader(
          meta,
          Lucene99HnswVectorsFormat.META_CODEC_NAME,
          version,
          state.segmentInfo.getId(),
          state.segmentSuffix);
      CodecUtil.writeIndexHeader(
          vectorIndex,
          Lucene99HnswVectorsFormat.VECTOR_INDEX_CODEC_NAME,
          version,
          state.segmentInfo.getId(),
          state.segmentSuffix);
      success = true;
    } finally {
      if (success == false) {
        IOUtils.closeWhileHandlingException(this);
      }
    }
  }

  @Override
  public KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException {
    FieldWriter<?> newField =
        FieldWriter.create(
            flatVectorWriter.getFlatVectorScorer(),
            flatVectorWriter.addField(fieldInfo),
            fieldInfo,
            M,
            beamWidth,
            segmentWriteState.infoStream,
            tinySegmentsThreshold);
    fields.add(newField);
    return newField;
  }

  @Override
  public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
    flatVectorWriter.flush(maxDoc, sortMap);
    for (FieldWriter<?> field : fields) {
      if (sortMap == null) {
        writeField(field);
      } else {
        writeSortingField(field, sortMap);
      }
    }
  }

  @Override
  public void finish() throws IOException {
    if (finished) {
      throw new IllegalStateException("already finished");
    }
    finished = true;
    if (flatWriterClosed == false) {
      flatVectorWriter.finish();
    }

    if (meta != null) {
      // write end of fields marker
      meta.writeInt(-1);
      CodecUtil.writeFooter(meta);
    }
    if (vectorIndex != null) {
      CodecUtil.writeFooter(vectorIndex);
    }
  }

  @Override
  public long ramBytesUsed() {
    long total = SHALLOW_RAM_BYTES_USED;
    for (FieldWriter<?> field : fields) {
      // the field tracks the delegate field usage
      total += field.ramBytesUsed();
    }
    return total;
  }

  private void writeField(FieldWriter<?> fieldData) throws IOException {
    // write graph
    long vectorIndexOffset = vectorIndex.getFilePointer();
    OnHeapHnswGraph graph = fieldData.getGraph();
    int[][] graphLevelNodeOffsets = writeGraph(graph);
    long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset;

    writeMeta(
        fieldData.fieldInfo,
        vectorIndexOffset,
        vectorIndexLength,
        fieldData.getDocsWithFieldSet().cardinality(),
        graph,
        graphLevelNodeOffsets);
  }

  private void writeSortingField(FieldWriter<?> fieldData, Sorter.DocMap sortMap)
      throws IOException {
    final int[] ordMap =
        new int[fieldData.getDocsWithFieldSet().cardinality()]; // new ord to old ord
    final int[] oldOrdMap =
        new int[fieldData.getDocsWithFieldSet().cardinality()]; // old ord to new ord

    mapOldOrdToNewOrd(fieldData.getDocsWithFieldSet(), sortMap, oldOrdMap, ordMap, null);
    // write graph
    long vectorIndexOffset = vectorIndex.getFilePointer();
    OnHeapHnswGraph graph = fieldData.getGraph();
    int[][] graphLevelNodeOffsets = graph == null ? new int[0][] : new int[graph.numLevels()][];
    HnswGraph mockGraph = reconstructAndWriteGraph(graph, ordMap, oldOrdMap, graphLevelNodeOffsets);
    long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset;

    writeMeta(
        fieldData.fieldInfo,
        vectorIndexOffset,
        vectorIndexLength,
        fieldData.getDocsWithFieldSet().cardinality(),
        mockGraph,
        graphLevelNodeOffsets);
  }

  /**
   * Reconstructs the graph given the old and new node ids.
   *
   * <p>Additionally, the graph node connections are written to the vectorIndex.
   *
   * @param graph The current on heap graph
   * @param newToOldMap the new node ids indexed to the old node ids
   * @param oldToNewMap the old node ids indexed to the new node ids
   * @param levelNodeOffsets where to place the new offsets for the nodes in the vector index.
   * @return The graph
   * @throws IOException if writing to vectorIndex fails
   */
  private HnswGraph reconstructAndWriteGraph(
      OnHeapHnswGraph graph, int[] newToOldMap, int[] oldToNewMap, int[][] levelNodeOffsets)
      throws IOException {
    if (graph == null) return null;

    List<int[]> nodesByLevel = new ArrayList<>(graph.numLevels());
    nodesByLevel.add(null);

    int maxOrd = graph.size();
    int[] scratch = new int[graph.maxConn() * 2];
    NodesIterator nodesOnLevel0 = graph.getNodesOnLevel(0);
    levelNodeOffsets[0] = new int[nodesOnLevel0.size()];
    while (nodesOnLevel0.hasNext()) {
      int node = nodesOnLevel0.nextInt();
      NeighborArray neighbors = graph.getNeighbors(0, newToOldMap[node]);
      long offset = vectorIndex.getFilePointer();
      reconstructAndWriteNeighbours(neighbors, oldToNewMap, scratch, maxOrd);
      levelNodeOffsets[0][node] = Math.toIntExact(vectorIndex.getFilePointer() - offset);
    }

    for (int level = 1; level < graph.numLevels(); level++) {
      NodesIterator nodesOnLevel = graph.getNodesOnLevel(level);
      int[] newNodes = new int[nodesOnLevel.size()];
      for (int n = 0; nodesOnLevel.hasNext(); n++) {
        newNodes[n] = oldToNewMap[nodesOnLevel.nextInt()];
      }
      Arrays.sort(newNodes);
      nodesByLevel.add(newNodes);
      levelNodeOffsets[level] = new int[newNodes.length];
      int nodeOffsetIndex = 0;
      for (int node : newNodes) {
        NeighborArray neighbors = graph.getNeighbors(level, newToOldMap[node]);
        long offset = vectorIndex.getFilePointer();
        reconstructAndWriteNeighbours(neighbors, oldToNewMap, scratch, maxOrd);
        levelNodeOffsets[level][nodeOffsetIndex++] =
            Math.toIntExact(vectorIndex.getFilePointer() - offset);
      }
    }
    return new HnswGraph() {
      @Override
      public int nextNeighbor() {
        throw new UnsupportedOperationException("Not supported on a mock graph");
      }

      @Override
      public void seek(int level, int target) {
        throw new UnsupportedOperationException("Not supported on a mock graph");
      }

      @Override
      public int size() {
        return graph.size();
      }

      @Override
      public int numLevels() {
        return graph.numLevels();
      }

      @Override
      public int maxConn() {
        return graph.maxConn();
      }

      @Override
      public int entryNode() {
        throw new UnsupportedOperationException("Not supported on a mock graph");
      }

      @Override
      public int neighborCount() {
        throw new UnsupportedOperationException("Not supported on a mock graph");
      }

      @Override
      public NodesIterator getNodesOnLevel(int level) {
        if (level == 0) {
          return graph.getNodesOnLevel(0);
        } else {
          return new ArrayNodesIterator(nodesByLevel.get(level));
        }
      }
    };
  }

  private void reconstructAndWriteNeighbours(
      NeighborArray neighbors, int[] oldToNewMap, int[] scratch, int maxOrd) throws IOException {
    int size = neighbors.size();
    // Destructively modify; it's ok we are discarding it after this
    int[] nnodes = neighbors.nodes();
    for (int i = 0; i < size; i++) {
      nnodes[i] = oldToNewMap[nnodes[i]];
    }
    Arrays.sort(nnodes, 0, size);
    int actualSize = 0;
    if (size > 0) {
      scratch[0] = nnodes[0];
      actualSize = 1;
    }
    // Now that we have sorted, do delta encoding to minimize the required bits to store the
    // information
    for (int i = 1; i < size; i++) {
      assert nnodes[i] < maxOrd : "node too large: " + nnodes[i] + ">=" + maxOrd;
      if (nnodes[i - 1] == nnodes[i]) {
        continue;
      }
      scratch[actualSize++] = nnodes[i] - nnodes[i - 1];
    }
    // Write the size after duplicates are removed
    vectorIndex.writeVInt(actualSize);
    if (version >= VERSION_GROUPVARINT) {
      vectorIndex.writeGroupVInts(scratch, actualSize);
    } else {
      for (int i = 0; i < actualSize; i++) {
        vectorIndex.writeVInt(scratch[i]);
      }
    }
  }

  @Override
  public IORunnable mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
    flatVectorWriter.mergeOneFlatVectorField(fieldInfo, mergeState);
    return () -> {
      // Lazily finish flat writer and open a reader for the written segment
      ensureFlatReaderOpen();
      // Get the vector values and scorer supplier from the written segment
      KnnVectorValues vectorValues =
          switch (fieldInfo.getVectorEncoding()) {
            case BYTE -> flatVectorsReader.getByteVectorValues(fieldInfo.name);
            case FLOAT32 -> flatVectorsReader.getFloatVectorValues(fieldInfo.name);
          };
      int totalVectorCount = vectorValues == null ? 0 : vectorValues.size();
      if (totalVectorCount > 0 && shouldCreateGraph(tinySegmentsThreshold, totalVectorCount)) {
        if (flatVectorsReader instanceof QuantizedVectorsReader quantizedVectorsReader
            && fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
          CloseableRandomVectorScorerSupplier scorerSupplier =
              quantizedVectorsReader.getRandomVectorScorerSupplierForMerge(
                  fieldInfo, segmentWriteState);
          try {
            buildAndWriteGraph(
                fieldInfo, mergeState, vectorValues, scorerSupplier, totalVectorCount);
          } catch (Throwable t) {
            IOUtils.closeWhileSuppressingExceptions(t, scorerSupplier);
            throw t;
          }
          IOUtils.close(scorerSupplier);
        } else {
          RandomVectorScorerSupplier scorerSupplier =
              flatVectorsReader
                  .getFlatVectorScorer(fieldInfo.getName())
                  .getRandomVectorScorerSupplier(
                      fieldInfo.getVectorSimilarityFunction(), vectorValues);
          buildAndWriteGraph(fieldInfo, mergeState, vectorValues, scorerSupplier, totalVectorCount);
        }
      } else {
        writeMeta(fieldInfo, vectorIndex.getFilePointer(), 0L, totalVectorCount, null, null);
      }
    };
  }

  private void buildAndWriteGraph(
      FieldInfo fieldInfo,
      MergeState mergeState,
      KnnVectorValues vectorValues,
      RandomVectorScorerSupplier scorerSupplier,
      int totalVectorCount)
      throws IOException {
    long vectorIndexOffset = vectorIndex.getFilePointer();
    OnHeapHnswGraph graph = null;
    int[][] vectorIndexNodeOffsets = null;
    HnswGraphMerger merger =
        createGraphMerger(
            fieldInfo,
            scorerSupplier,
            mergeState.intraMergeTaskExecutor == null
                ? null
                : new TaskExecutor(mergeState.intraMergeTaskExecutor),
            numMergeWorkers);
    for (int i = 0; i < mergeState.liveDocs.length; i++) {
      if (hasVectorValues(mergeState.fieldInfos[i], fieldInfo.name)) {
        merger.addReader(
            mergeState.knnVectorsReaders[i], mergeState.docMaps[i], mergeState.liveDocs[i]);
      }
    }
    graph = merger.merge(vectorValues, segmentWriteState.infoStream, totalVectorCount);
    vectorIndexNodeOffsets = writeGraph(graph);
    long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset;
    writeMeta(
        fieldInfo,
        vectorIndexOffset,
        vectorIndexLength,
        totalVectorCount,
        graph,
        vectorIndexNodeOffsets);
  }

  private void ensureFlatReaderOpen() throws IOException {
    if (flatVectorsReader == null) {
      flatVectorWriter.finish();
      flatVectorWriter.close();
      flatWriterClosed = true;
      SegmentReadState readState =
          new SegmentReadState(
              segmentWriteState.directory,
              segmentWriteState.segmentInfo,
              segmentWriteState.fieldInfos,
              segmentWriteState.context,
              segmentWriteState.segmentSuffix);
      flatVectorsReader = flatVectorsFormat.fieldsReader(readState);
    }
  }

  /**
   * @param graph Write the graph in a compressed format
   * @return The non-cumulative offsets for the nodes. Should be used to create cumulative offsets.
   * @throws IOException if writing to vectorIndex fails
   */
  private int[][] writeGraph(OnHeapHnswGraph graph) throws IOException {
    if (graph == null) return new int[0][0];
    // write vectors' neighbours on each level into the vectorIndex file
    int countOnLevel0 = graph.size();
    int[][] offsets = new int[graph.numLevels()][];
    int[] scratch = new int[graph.maxConn() * 2];
    for (int level = 0; level < graph.numLevels(); level++) {
      NodesIterator sortedNodes = graph.getSortedNodes(level);
      offsets[level] = new int[sortedNodes.size()];
      int nodeOffsetId = 0;
      while (sortedNodes.hasNext()) {
        NeighborArray neighbors = graph.getNeighbors(level, sortedNodes.next());
        int size = neighbors.size();
        // Write size in VInt as the neighbors list is typically small
        long offsetStart = vectorIndex.getFilePointer();
        int[] nnodes = neighbors.nodes();
        Arrays.sort(nnodes, 0, size);
        // Now that we have sorted, do delta encoding to minimize the required bits to store the
        // information
        int actualSize = 0;
        if (size > 0) {
          scratch[0] = nnodes[0];
          actualSize = 1;
        }
        for (int i = 1; i < size; i++) {
          assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0;
          if (nnodes[i - 1] == nnodes[i]) {
            continue;
          }
          scratch[actualSize++] = nnodes[i] - nnodes[i - 1];
        }
        // Write the size after duplicates are removed
        vectorIndex.writeVInt(actualSize);
        if (version >= VERSION_GROUPVARINT) {
          vectorIndex.writeGroupVInts(scratch, actualSize);
        } else {
          for (int i = 0; i < actualSize; i++) {
            vectorIndex.writeVInt(scratch[i]);
          }
        }

        offsets[level][nodeOffsetId++] =
            Math.toIntExact(vectorIndex.getFilePointer() - offsetStart);
      }
    }
    return offsets;
  }

  private void writeMeta(
      FieldInfo field,
      long vectorIndexOffset,
      long vectorIndexLength,
      int count,
      HnswGraph graph,
      int[][] graphLevelNodeOffsets)
      throws IOException {
    meta.writeInt(field.number);
    meta.writeInt(field.getVectorEncoding().ordinal());
    meta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction()));
    meta.writeVLong(vectorIndexOffset);
    meta.writeVLong(vectorIndexLength);
    meta.writeVInt(field.getVectorDimension());
    meta.writeInt(count);
    meta.writeVInt(M);
    // write graph nodes on each level
    if (graph == null) {
      meta.writeVInt(0);
    } else {
      meta.writeVInt(graph.numLevels());
      long valueCount = 0;
      for (int level = 0; level < graph.numLevels(); level++) {
        NodesIterator nodesOnLevel = graph.getNodesOnLevel(level);
        valueCount += nodesOnLevel.size();
        if (level > 0) {
          int[] nol = new int[nodesOnLevel.size()];
          int numberConsumed = nodesOnLevel.consume(nol);
          Arrays.sort(nol);
          assert numberConsumed == nodesOnLevel.size();
          meta.writeVInt(nol.length); // number of nodes on a level
          for (int i = nodesOnLevel.size() - 1; i > 0; --i) {
            nol[i] -= nol[i - 1];
          }
          for (int n : nol) {
            assert n >= 0 : "delta encoding for nodes failed; expected nodes to be sorted";
            meta.writeVInt(n);
          }
        } else {
          assert nodesOnLevel.size() == count : "Level 0 expects to have all nodes";
        }
      }
      long start = vectorIndex.getFilePointer();
      meta.writeLong(start);
      meta.writeVInt(DIRECT_MONOTONIC_BLOCK_SHIFT);
      final DirectMonotonicWriter memoryOffsetsWriter =
          DirectMonotonicWriter.getInstance(
              meta, vectorIndex, valueCount, DIRECT_MONOTONIC_BLOCK_SHIFT);
      long cumulativeOffsetSum = 0;
      for (int[] levelOffsets : graphLevelNodeOffsets) {
        for (int v : levelOffsets) {
          memoryOffsetsWriter.add(cumulativeOffsetSum);
          cumulativeOffsetSum += v;
        }
      }
      memoryOffsetsWriter.finish();
      meta.writeLong(vectorIndex.getFilePointer() - start);
    }
  }

  private HnswGraphMerger createGraphMerger(
      FieldInfo fieldInfo,
      RandomVectorScorerSupplier scorerSupplier,
      TaskExecutor parallelMergeTaskExecutor,
      int numParallelMergeWorkers) {
    if (mergeExec != null) {
      return new ConcurrentHnswMerger(
          fieldInfo, scorerSupplier, M, beamWidth, mergeExec, numMergeWorkers);
    }
    if (parallelMergeTaskExecutor != null && numParallelMergeWorkers > 1) {
      return new ConcurrentHnswMerger(
          fieldInfo,
          scorerSupplier,
          M,
          beamWidth,
          parallelMergeTaskExecutor,
          numParallelMergeWorkers);
    }
    return new IncrementalHnswGraphMerger(fieldInfo, scorerSupplier, M, beamWidth);
  }

  @Override
  public void close() throws IOException {
    if (flatWriterClosed) {
      IOUtils.close(meta, vectorIndex, flatVectorsReader);
    } else {
      IOUtils.close(meta, vectorIndex, flatVectorWriter, flatVectorsReader);
    }
  }

  static int distFuncToOrd(VectorSimilarityFunction func) {
    for (int i = 0; i < SIMILARITY_FUNCTIONS.size(); i++) {
      if (SIMILARITY_FUNCTIONS.get(i).equals(func)) {
        return (byte) i;
      }
    }
    throw new IllegalArgumentException("invalid distance function: " + func);
  }

  private static boolean shouldCreateGraph(int k, int numNodes) {
    if (k <= 0) {
      return true;
    }
    int expectedVisitedNodes = expectedVisitedNodes(k, numNodes);
    return numNodes > expectedVisitedNodes && expectedVisitedNodes > 0;
  }

  private static class FieldWriter<T> extends KnnFieldVectorsWriter<T> {

    private static final long SHALLOW_SIZE =
        RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class);

    private final FieldInfo fieldInfo;
    private HnswGraphBuilder hnswGraphBuilder; // only created when needed
    private int lastDocID = -1;
    private int node = 0;
    private final FlatFieldVectorsWriter<T> flatFieldVectorsWriter;
    private final int graphThreshold;
    private final int M;
    private final int beamWidth;
    private final InfoStream infoStream;
    private final RandomVectorScorerSupplier scorerSupplier;

    @SuppressWarnings("unchecked")
    static FieldWriter<?> create(
        FlatVectorsScorer scorer,
        FlatFieldVectorsWriter<?> flatFieldVectorsWriter,
        FieldInfo fieldInfo,
        int M,
        int beamWidth,
        InfoStream infoStream,
        int tinySegmentsThreshold)
        throws IOException {
      return switch (fieldInfo.getVectorEncoding()) {
        case BYTE ->
            new FieldWriter<>(
                scorer,
                (FlatFieldVectorsWriter<byte[]>) flatFieldVectorsWriter,
                fieldInfo,
                M,
                beamWidth,
                infoStream,
                tinySegmentsThreshold);
        case FLOAT32 ->
            new FieldWriter<>(
                scorer,
                (FlatFieldVectorsWriter<float[]>) flatFieldVectorsWriter,
                fieldInfo,
                M,
                beamWidth,
                infoStream,
                tinySegmentsThreshold);
      };
    }

    @SuppressWarnings("unchecked")
    FieldWriter(
        FlatVectorsScorer scorer,
        FlatFieldVectorsWriter<T> flatFieldVectorsWriter,
        FieldInfo fieldInfo,
        int M,
        int beamWidth,
        InfoStream infoStream,
        int tinySegmentsThreshold)
        throws IOException {
      this.fieldInfo = fieldInfo;
      this.M = M;
      this.beamWidth = beamWidth;
      this.infoStream = infoStream;
      this.flatFieldVectorsWriter = Objects.requireNonNull(flatFieldVectorsWriter);
      this.graphThreshold = tinySegmentsThreshold;
      this.scorerSupplier =
          scorer.getRandomVectorScorerSupplier(
              fieldInfo.getVectorSimilarityFunction(),
              flatFieldVectorsWriter.asKnnVectorValues(
                  fieldInfo.getVectorEncoding(), fieldInfo.getVectorDimension()));

      if (graphThreshold <= 0) {
        // Initialize graph builder if optimization is disabled
        initializeGraphBuilder();
      }
    }

    private void initializeGraphBuilder() throws IOException {
      if (hnswGraphBuilder != null) {
        return;
      }
      this.hnswGraphBuilder =
          HnswGraphBuilder.create(scorerSupplier, M, beamWidth, HnswGraphBuilder.randSeed);
      this.hnswGraphBuilder.setInfoStream(infoStream);
    }

    private void replayBufferedVectors() throws IOException {
      for (int i = 0; i < flatFieldVectorsWriter.getVectors().size(); i++) {
        hnswGraphBuilder.addGraphNode(i);
      }
    }

    @Override
    public void addValue(int docID, T vectorValue) throws IOException {
      if (docID == lastDocID) {
        throw new IllegalArgumentException(
            "VectorValuesField \""
                + fieldInfo.name
                + "\" appears more than once in this document (only one value is allowed per field)");
      }
      flatFieldVectorsWriter.addValue(docID, vectorValue);
      // Check if we need to initialize graph builder for tiny segment optimization
      if (hnswGraphBuilder == null && shouldCreateGraph(graphThreshold, node + 1)) {
        initializeGraphBuilder();
        replayBufferedVectors();
      } else if (hnswGraphBuilder != null) {
        // Graph builder is active, add to graph
        hnswGraphBuilder.addGraphNode(node);
      }
      node++;
      lastDocID = docID;
    }

    public DocsWithFieldSet getDocsWithFieldSet() {
      return flatFieldVectorsWriter.getDocsWithFieldSet();
    }

    @Override
    public T copyValue(T vectorValue) {
      throw new UnsupportedOperationException();
    }

    OnHeapHnswGraph getGraph() throws IOException {
      assert flatFieldVectorsWriter.isFinished();
      if (hnswGraphBuilder != null && node > 0) {
        return hnswGraphBuilder.getCompletedGraph();
      } else {
        // No graph
        return null;
      }
    }

    @Override
    public long ramBytesUsed() {
      long total = SHALLOW_SIZE + flatFieldVectorsWriter.ramBytesUsed();
      if (hnswGraphBuilder != null) {
        total += hnswGraphBuilder.getGraph().ramBytesUsed();
      }
      return total;
    }
  }
}
