package engine

import (
	"context"
	"math"
	"strings"
	"sync"

	"github.com/SimonWaldherr/tinySQL/internal/storage"
)

const (
	vecIndexFlat = "flat"
	vecIndexIVF  = "ivf"
	vecIndexHNSW = "hnsw"

	vecIVFKMeansIters = 3

	vecHNSWM              = 12
	vecHNSWEfConstruction = 48
	vecHNSWMaxLevel       = 8

	// vecIndexCacheMaxEntries bounds the IVF and HNSW caches — see the
	// vecColumnCacheMaxEntries comment in vector_search.go for rationale.
	// Lower than the column-cache cap because index entries are larger
	// (centroids/graph) and each also pins its *storage.Table.
	vecIndexCacheMaxEntries = 64
)

type vecIndexCacheKey struct {
	tenant string
	table  string
	colIdx int
	metric string
}

// vecIndexBuildCall coalesces concurrent cold ANN index builds for one cache
// key — see vecIVFBuilds/vecHNSWBuilds and vecColumnBuildCall (vector_search.go)
// for the shared pattern.
type vecIndexBuildCall struct{ done chan struct{} }

type vecIVFIndex struct {
	table         *storage.Table
	version       int
	metric        string
	dims          int
	centroids     [][]float64
	centroidNorms []float64
	lists         [][]int
}

type vecHNSWIndex struct {
	table   *storage.Table
	version int
	// structVersion is the table's StructVersion() as of the last time this
	// graph was fully built or incrementally extended. Equal to the table's
	// current StructVersion() proves every row this graph already indexes
	// (rows [0, len(levels))) is unchanged, so the only possible reason
	// idx.version can be behind table.Version is that rows were appended —
	// see canExtendVecHNSWIndex.
	structVersion int
	metric        string
	dims          int
	entry         int
	maxLevel      int
	levels        []int
	neighbors     [][][]int
	// mu guards every field above against extendVecHNSWIndex mutating an
	// already-published graph (see getVecHNSWIndex) while a concurrent
	// search() call is reading it. A freshly built graph (buildVecHNSWIndex)
	// needs no locking: it is not reachable through the cache, and therefore
	// not visible to any other goroutine, until after it is fully built.
	mu sync.RWMutex
}

type vecMinScoredHeap []vecScoredRow

func (h vecMinScoredHeap) Len() int { return len(h) }
func (h vecMinScoredHeap) Less(i, j int) bool {
	return vecScoredRowLess(h[i], h[j])
}
func (h vecMinScoredHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

// vecMinHeapPush/Pop mirror vecScoredHeapPush/Pop (see vector_search.go) for
// the candidates frontier — direct concrete-type push/pop instead of
// heap.Push/heap.Pop, which box every vecScoredRow into an `any` and
// allocate. searchLayer calls this once per neighbor considered during graph
// traversal, so this was the largest remaining allocation source in HNSW
// build/search.
func vecMinHeapPush(h *vecMinScoredHeap, v vecScoredRow) {
	*h = append(*h, v)
	vecMinHeapUp(*h, len(*h)-1)
}

func vecMinHeapPop(h *vecMinScoredHeap) vecScoredRow {
	old := *h
	n := len(old) - 1
	old.Swap(0, n)
	vecMinHeapDown(old[:n], 0)
	v := old[n]
	*h = old[:n]
	return v
}

func vecMinHeapUp(h vecMinScoredHeap, j int) {
	for {
		i := (j - 1) / 2
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}
}

func vecMinHeapDown(h vecMinScoredHeap, i0 int) {
	n := len(h)
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 {
			break
		}
		j := j1
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
			j = j2
		}
		if !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		i = j
	}
}

var (
	vecIVFCacheMu sync.RWMutex
	vecIVFCache   = make(map[vecIndexCacheKey]*vecIVFIndex)
	// vecIVFBuilds coalesces concurrent cold IVF builds for the same
	// (tenant, table, colIdx, metric) key, mirroring vecSearchColumnBuilds in
	// vector_search.go. Keying the in-flight-build map by the same cache key
	// (rather than one global mutex) lets concurrent builds for DIFFERENT
	// keys proceed in parallel — e.g. VEC_WARM warming several tables at
	// startup takes max(build_i) instead of sum(build_i) — while concurrent
	// requests for the SAME key still coalesce onto a single build.
	vecIVFBuilds = make(map[vecIndexCacheKey]*vecIndexBuildCall)

	vecHNSWCacheMu sync.RWMutex
	vecHNSWCache   = make(map[vecIndexCacheKey]*vecHNSWIndex)
	// vecHNSWBuilds is vecIVFBuilds' counterpart for the HNSW cache.
	vecHNSWBuilds = make(map[vecIndexCacheKey]*vecIndexBuildCall)

	// vecVisitedPool recycles the per-search visited bitmaps so HNSW queries
	// on large tables do not allocate len(rows) bytes per call.
	vecVisitedPool sync.Pool
)

// vecHNSWScratch holds the reusable candidate/result heaps and visited-touch
// list for one HNSW traversal (one search() call, or the whole sequential
// build loop). searchLayer runs once per graph layer visited during a
// traversal, and previously allocated a fresh pair of heaps plus a touched-
// node slice on every one of those calls — for a 12k-row build that added up
// to tens of thousands of small allocations and made container/heap's
// interface-dispatched Push/Pop the dominant CPU cost. Reusing the backing
// arrays (reset to length 0, capacity retained) across calls, and across
// pooled instances between queries, converges to zero steady-state
// allocations for the traversal's own bookkeeping.
type vecHNSWScratch struct {
	candidates vecMinScoredHeap
	results    vecScoredHeap
	touched    []int
}

var vecHNSWScratchPool = sync.Pool{
	New: func() any { return new(vecHNSWScratch) },
}

func acquireHNSWScratch() *vecHNSWScratch {
	return vecHNSWScratchPool.Get().(*vecHNSWScratch)
}

func releaseHNSWScratch(s *vecHNSWScratch) {
	vecHNSWScratchPool.Put(s)
}

func acquireVisited(n int) []bool {
	if v, ok := vecVisitedPool.Get().([]bool); ok && cap(v) >= n {
		v = v[:n]
		clear(v)
		return v
	}
	// Grow with slack (like append's own doubling growth) rather than
	// exactly n: a caller whose required size creeps up by a little on
	// every call — e.g. extendVecHNSWIndex, called once per single-row
	// INSERT, with n following the table's row count up by one each time —
	// would otherwise force a fresh allocation on every single call despite
	// the pool existing, since the previous release's capacity is always
	// exactly one short of what the next acquire needs. The extra
	// capacity is never visible to callers (every acquireVisited caller
	// only ever sees the v[:n] returned here), so this is purely an
	// internal amortization, not a behavior change.
	return make([]bool, n, n+n/2+64)
}

func releaseVisited(v []bool) {
	if cap(v) > 0 {
		vecVisitedPool.Put(v[:0]) //nolint:staticcheck // slice header, not pointer
	}
}

func normalizeVecIndexMode(mode string) string {
	switch strings.ToLower(strings.TrimSpace(mode)) {
	case "", "flat", "exact", "scan":
		return vecIndexFlat
	case "ivf":
		return vecIndexIVF
	case "hnsw":
		return vecIndexHNSW
	default:
		return ""
	}
}

// vecSearchTopKWithIndex is the single choke point every index mode's
// search funnels through before a result ever reaches a caller. Internally,
// flat/IVF/HNSW all rank candidates using vectorRankingDistance (skipping
// l2's sqrt until it's actually needed); finalizeVecScoredRows converts the
// (at most k) surviving rows back to real distances exactly once here,
// rather than every candidate the scan/traversal considered paying for it.
func vecSearchTopKWithIndex(
	ctx context.Context,
	tenant string,
	table *storage.Table,
	colIdx int,
	args vecSearchArgs,
	queryLen int,
	queryNorm float64,
	cache vecSearchColumnCacheEntry,
	distFn vecDistanceFunc,
) ([]vecScoredRow, error) {
	var rows []vecScoredRow
	var err error
	switch args.indexMode {
	case vecIndexFlat:
		rows, err = vecSearchTopK(ctx, table.Rows, queryLen, args.k, cache, distFn)
	case vecIndexIVF:
		var idx *vecIVFIndex
		idx, err = getVecIVFIndex(ctx, tenant, table, colIdx, args.metric, queryLen, cache)
		if err == nil {
			rows, err = idx.search(ctx, args.queryVec, queryNorm, args.k, cache)
		}
	case vecIndexHNSW:
		var idx *vecHNSWIndex
		idx, err = getVecHNSWIndex(ctx, tenant, table, colIdx, args.metric, queryLen, cache)
		if err == nil {
			rows, err = idx.search(ctx, args.queryVec, queryNorm, args.k, cache)
		}
	default:
		rows, err = vecSearchTopK(ctx, table.Rows, queryLen, args.k, cache, distFn)
	}
	if err != nil {
		return nil, err
	}
	finalizeVecScoredRows(args.metric, rows)
	return rows, nil
}

// finalizeVecScoredRows converts every row's ranking-only distance (see
// vectorRankingDistance) into the real distance vectorDistance would have
// returned, in place. Cheap to call unconditionally: for every metric
// except l2 vectorFinalizeDistance is already a no-op, and rows is bounded
// to k (or fewer) entries regardless of how many candidates the scan or
// graph traversal that produced it considered.
func finalizeVecScoredRows(metric string, rows []vecScoredRow) {
	if metric != "l2" {
		return
	}
	for i := range rows {
		rows[i].distance = vectorFinalizeDistance(metric, rows[i].distance)
	}
}

func getVecIVFIndex(ctx context.Context, tenant string, table *storage.Table, colIdx int, metric string, dims int, cache vecSearchColumnCacheEntry) (*vecIVFIndex, error) {
	key := vecIndexCacheKey{tenant: tenant, table: table.Name, colIdx: colIdx, metric: metric}

	for {
		vecIVFCacheMu.RLock()
		if idx := vecIVFCache[key]; idx != nil && idx.table == table && idx.version == table.Version && idx.dims == dims {
			vecIVFCacheMu.RUnlock()
			return idx, nil
		}
		vecIVFCacheMu.RUnlock()

		vecIVFCacheMu.Lock()
		if idx := vecIVFCache[key]; idx != nil && idx.table == table && idx.version == table.Version && idx.dims == dims {
			vecIVFCacheMu.Unlock()
			return idx, nil
		}
		if call := vecIVFBuilds[key]; call != nil {
			vecIVFCacheMu.Unlock()
			<-call.done
			continue
		}
		call := &vecIndexBuildCall{done: make(chan struct{})}
		vecIVFBuilds[key] = call
		vecIVFCacheMu.Unlock()

		idx, err := buildVecIVFIndex(ctx, table, metric, dims, cache)

		vecIVFCacheMu.Lock()
		delete(vecIVFBuilds, key)
		if err == nil {
			if _, exists := vecIVFCache[key]; !exists {
				evictOverCap(vecIVFCache, vecIndexCacheMaxEntries)
			}
			vecIVFCache[key] = idx
		}
		close(call.done)
		vecIVFCacheMu.Unlock()
		return idx, err
	}
}

func buildVecIVFIndex(ctx context.Context, table *storage.Table, metric string, dims int, cache vecSearchColumnCacheEntry) (*vecIVFIndex, error) {
	rows := validVectorRows(cache, dims)
	nlist := chooseIVFListCount(len(rows))
	idx := &vecIVFIndex{table: table, version: table.Version, metric: metric, dims: dims}
	if len(rows) == 0 || nlist == 0 {
		return idx, nil
	}

	// Seed centroids by taking nlist rows spread evenly across the (already
	// arbitrary-ordered) row set, rather than the first nlist rows — cheap
	// and avoids every centroid starting near-identical if the table happens
	// to be sorted or clustered by insertion order.
	idx.centroids = make([][]float64, nlist)
	for i := range idx.centroids {
		src := cache.vectors[rows[(i*len(rows))/nlist]]
		idx.centroids[i] = append([]float64(nil), src...)
	}
	assignments := make([]int, len(rows))
	// Reuse the accumulation buffers across k-means iterations instead of
	// reallocating nlist*dims floats per pass.
	sums := make([]float64, nlist*dims)
	counts := make([]int, nlist)
	for iter := 0; iter < vecIVFKMeansIters; iter++ {
		if err := checkCtx(ctx); err != nil {
			return nil, err
		}
		idx.centroidNorms = centroidNormsFor(metric, idx.centroids)
		clear(sums)
		clear(counts)
		for i, rowIdx := range rows {
			if i&1023 == 0 {
				if err := checkCtx(ctx); err != nil {
					return nil, err
				}
			}
			c := nearestCentroid(metric, cache.vectors[rowIdx], rowNormFor(metric, cache, rowIdx), idx.centroids, idx.centroidNorms)
			assignments[i] = c
			counts[c]++
			base := c * dims
			vectorAccumulateUnrolled(sums[base:base+dims], cache.vectors[rowIdx])
		}
		for c := range idx.centroids {
			// A centroid can end up with zero assigned rows this iteration
			// (every row happened to be closer to some other centroid) —
			// a normal k-means edge case, more likely with unlucky seeding
			// or nlist close to the number of distinct clusters actually
			// present. Leave it at its previous position rather than
			// dividing by zero; it may pick up members in a later
			// iteration as other centroids move, or simply stay unused.
			if counts[c] == 0 {
				continue
			}
			inv := 1.0 / float64(counts[c])
			base := c * dims
			for d := 0; d < dims; d++ {
				idx.centroids[c][d] = sums[base+d] * inv
			}
		}
	}

	idx.centroidNorms = centroidNorms(idx.centroids)
	idx.lists = make([][]int, nlist)
	for i, rowIdx := range rows {
		c := assignments[i]
		idx.lists[c] = append(idx.lists[c], rowIdx)
	}
	return idx, nil
}

func (idx *vecIVFIndex) search(ctx context.Context, query []float64, queryNorm float64, k int, cache vecSearchColumnCacheEntry) ([]vecScoredRow, error) {
	if len(idx.centroids) == 0 {
		return nil, nil
	}
	probes := chooseIVFNProbe(len(idx.centroids), k)
	centroidHeap := &vecScoredHeap{}
	for i, c := range idx.centroids {
		// Selecting which centroids to probe, not a value ever exposed to a
		// caller — ranking-only distance is enough (see rowDistance below
		// for the row-level equivalent).
		dist, ok := vectorRankingDistance(idx.metric, c, query, idx.centroidNorms[i], queryNorm)
		if !ok {
			continue
		}
		pushTopK(centroidHeap, i, dist, probes)
	}
	bestCentroids := topKFromHeap(centroidHeap, probes)

	resultHeap := &vecScoredHeap{}
	for _, c := range bestCentroids {
		list := idx.lists[c.rowIdx]
		for i, rowIdx := range list {
			if i&1023 == 0 {
				if err := checkCtx(ctx); err != nil {
					return nil, err
				}
			}
			dist, ok := rowDistance(idx.metric, query, queryNorm, cache, rowIdx)
			if ok {
				pushTopK(resultHeap, rowIdx, dist, k)
			}
		}
	}
	return topKFromHeap(resultHeap, k), nil
}

func getVecHNSWIndex(ctx context.Context, tenant string, table *storage.Table, colIdx int, metric string, dims int, cache vecSearchColumnCacheEntry) (*vecHNSWIndex, error) {
	key := vecIndexCacheKey{tenant: tenant, table: table.Name, colIdx: colIdx, metric: metric}

	for {
		vecHNSWCacheMu.RLock()
		if idx := vecHNSWCache[key]; idx != nil && idx.table == table && idx.version == table.Version && idx.dims == dims {
			vecHNSWCacheMu.RUnlock()
			return idx, nil
		}
		vecHNSWCacheMu.RUnlock()

		vecHNSWCacheMu.Lock()
		idx := vecHNSWCache[key]
		if idx != nil && idx.table == table && idx.version == table.Version && idx.dims == dims {
			vecHNSWCacheMu.Unlock()
			return idx, nil
		}
		if call := vecHNSWBuilds[key]; call != nil {
			vecHNSWCacheMu.Unlock()
			<-call.done
			continue
		}

		// A cached graph that is stale only because rows were appended since
		// it was built or last extended (no UPDATE, DELETE, or schema change
		// happened in between — see canExtendVecHNSWIndex) can grow in place
		// instead of being discarded and rebuilt from scratch. This mirrors
		// the incremental-append/full-rebuild-on-anything-else split already
		// used for constraintIndexes (exec_dml_insert.go): appends are cheap
		// to absorb, everything else re-derives the whole structure.
		if canExtendVecHNSWIndex(idx, table, dims) {
			call := &vecIndexBuildCall{done: make(chan struct{})}
			vecHNSWBuilds[key] = call
			vecHNSWCacheMu.Unlock()

			idx.mu.Lock()
			err := extendVecHNSWIndex(ctx, idx, cache)
			if err == nil {
				idx.version = table.Version
				idx.structVersion = table.StructVersion()
			}
			idx.mu.Unlock()

			vecHNSWCacheMu.Lock()
			delete(vecHNSWBuilds, key)
			close(call.done)
			vecHNSWCacheMu.Unlock()
			if err != nil {
				return nil, err
			}
			return idx, nil
		}

		call := &vecIndexBuildCall{done: make(chan struct{})}
		vecHNSWBuilds[key] = call
		vecHNSWCacheMu.Unlock()

		newIdx, err := buildVecHNSWIndex(ctx, table, metric, dims, cache)

		vecHNSWCacheMu.Lock()
		delete(vecHNSWBuilds, key)
		if err == nil {
			if _, exists := vecHNSWCache[key]; !exists {
				evictOverCap(vecHNSWCache, vecIndexCacheMaxEntries)
			}
			vecHNSWCache[key] = newIdx
		}
		close(call.done)
		vecHNSWCacheMu.Unlock()
		return newIdx, err
	}
}

// canExtendVecHNSWIndex reports whether idx can grow in place to cover
// table's current rows instead of being rebuilt from scratch. This is safe
// exactly when table.StructVersion() has not moved since idx was last built
// or extended: that counter only advances on UPDATE, DELETE, or a schema
// change (see Table.noteStructuralChange), never on INSERT, so an unchanged
// value proves every row idx already indexes still holds exactly what it did
// then — the table can only have grown by appending. The row-count check is
// a second, independent guard against the same conclusion reached unsafely
// (e.g. a table rolled back to fewer rows than idx has already indexed,
// which would otherwise still show an unchanged StructVersion).
func canExtendVecHNSWIndex(idx *vecHNSWIndex, table *storage.Table, dims int) bool {
	return idx != nil &&
		idx.table == table &&
		idx.dims == dims &&
		idx.structVersion == table.StructVersion() &&
		len(idx.levels) <= len(table.Rows)
}

func buildVecHNSWIndex(ctx context.Context, table *storage.Table, metric string, dims int, cache vecSearchColumnCacheEntry) (*vecHNSWIndex, error) {
	idx := &vecHNSWIndex{
		table:         table,
		version:       table.Version,
		structVersion: table.StructVersion(),
		metric:        metric,
		dims:          dims,
		entry:         -1,
		maxLevel:      -1,
		levels:        make([]int, len(cache.vectors)),
		neighbors:     make([][][]int, len(cache.vectors)),
	}

	visited := make([]bool, len(cache.vectors))
	// The build loop is strictly sequential (single goroutine inserting one
	// row at a time), so a single scratch instance can be reused across all
	// insertions instead of round-tripping through the pool per row.
	scratch := acquireHNSWScratch()
	defer releaseHNSWScratch(scratch)
	for rowIdx := range cache.vectors {
		if rowIdx&1023 == 0 {
			if err := checkCtx(ctx); err != nil {
				return nil, err
			}
		}
		if !validCacheRow(cache, rowIdx, dims) {
			continue
		}
		idx.insertHNSWNode(rowIdx, cache, visited, scratch)
	}
	return idx, nil
}

// extendVecHNSWIndex grows an already-built HNSW graph in place by inserting
// every row appended to the table since idx was last built or extended,
// using insertHNSWNode — the exact same per-row insert logic
// buildVecHNSWIndex uses for a from-scratch build. The resulting graph is
// therefore structurally identical to one built from scratch on the same
// final data (same neighbor selection, same pruning, same entry-point
// evolution), not merely "a working graph".
//
// Callers must already have established, via canExtendVecHNSWIndex, that
// only appends — never an UPDATE, DELETE, or schema change — happened to
// the table since idx.version; this function does not re-check that and
// trusts idx's existing rows [0, len(idx.levels)) to still be exactly what
// they were.
//
// idx is the same object already published in vecHNSWCache, so the caller
// must hold idx.mu for writing for the duration of this call — a concurrent
// search() call may be reading idx through a reference obtained earlier.
//
// If ctx is cancelled partway through, idx.levels/idx.neighbors are
// truncated back to the last row fully inserted (every link created so far
// only ever points at rows below that point, since a node can only link to
// something already reachable in the graph) so idx is left in exactly the
// state a shorter, successful extend would have produced — never partially
// linked. idx.version/structVersion are left untouched on error, so the next
// getVecHNSWIndex call sees it as still eligible for (and resumes) an
// incremental extend from that point, rather than paying for a full rebuild.
func extendVecHNSWIndex(ctx context.Context, idx *vecHNSWIndex, cache vecSearchColumnCacheEntry) error {
	oldRows := len(idx.levels)
	newRows := len(cache.vectors)
	if newRows <= oldRows {
		return nil
	}
	idx.levels = append(idx.levels, make([]int, newRows-oldRows)...)
	idx.neighbors = append(idx.neighbors, make([][][]int, newRows-oldRows)...)

	// Sized to the final row count up front, exactly like buildVecHNSWIndex's
	// visited array, since traversal while inserting a new row can touch any
	// row already linked into the graph, old or newly inserted earlier in
	// this same call. Pooled (see acquireVisited/search) rather than a plain
	// make(): a from-scratch build pays one allocation for the whole build,
	// but every single-row INSERT that reaches this function would otherwise
	// pay a fresh O(existing rows) allocation of its own on every call — an
	// allocation cost that grows with table size defeating the entire point
	// of extending in place instead of rebuilding.
	visited := acquireVisited(newRows)
	defer releaseVisited(visited)
	scratch := acquireHNSWScratch()
	defer releaseHNSWScratch(scratch)
	for rowIdx := oldRows; rowIdx < newRows; rowIdx++ {
		if rowIdx&1023 == 0 {
			if err := checkCtx(ctx); err != nil {
				idx.levels = idx.levels[:rowIdx]
				idx.neighbors = idx.neighbors[:rowIdx]
				return err
			}
		}
		if !validCacheRow(cache, rowIdx, idx.dims) {
			continue
		}
		idx.insertHNSWNode(rowIdx, cache, visited, scratch)
	}
	return nil
}

// insertHNSWNode inserts rowIdx into idx's graph. It assumes idx.levels and
// idx.neighbors already have a slot allocated for every index up to rowIdx
// (as both buildVecHNSWIndex and extendVecHNSWIndex arrange before calling
// this), and that rows below rowIdx are already fully inserted. Shared by
// both, so a from-scratch build and an incremental extend grow the exact
// same graph structure via the exact same logic.
func (idx *vecHNSWIndex) insertHNSWNode(rowIdx int, cache vecSearchColumnCacheEntry, visited []bool, scratch *vecHNSWScratch) {
	level := hnswLevel(rowIdx)
	idx.levels[rowIdx] = level
	idx.neighbors[rowIdx] = make([][]int, level+1)
	if idx.entry < 0 {
		idx.entry = rowIdx
		idx.maxLevel = level
		return
	}

	current := idx.entry
	query := cache.vectors[rowIdx]
	queryNorm := rowNormFor(idx.metric, cache, rowIdx)
	for layer := idx.maxLevel; layer > level; layer-- {
		best := idx.searchLayer(query, queryNorm, current, 1, layer, cache, visited, scratch)
		if len(best) > 0 {
			current = best[0].rowIdx
		}
	}
	upper := level
	if idx.maxLevel < upper {
		upper = idx.maxLevel
	}
	for layer := upper; layer >= 0; layer-- {
		candidates := idx.searchLayer(query, queryNorm, current, vecHNSWEfConstruction, layer, cache, visited, scratch)
		selected := selectHNSWNeighbors(candidates, vecHNSWM)
		for _, nb := range selected {
			idx.addHNSWLink(rowIdx, nb.rowIdx, layer, cache)
		}
		if len(selected) > 0 {
			current = selected[0].rowIdx
		}
	}
	if level > idx.maxLevel {
		idx.entry = rowIdx
		idx.maxLevel = level
	}
}

func (idx *vecHNSWIndex) search(ctx context.Context, query []float64, queryNorm float64, k int, cache vecSearchColumnCacheEntry) ([]vecScoredRow, error) {
	// extendVecHNSWIndex can mutate this same, already-published object in
	// place (see getVecHNSWIndex); a plain full rebuild never touches an
	// object once published, so this read lock costs nothing extra in that
	// (still far more common) case beyond an uncontended RWMutex RLock.
	idx.mu.RLock()
	defer idx.mu.RUnlock()
	if idx.entry < 0 {
		return nil, nil
	}
	current := idx.entry
	visited := acquireVisited(len(cache.vectors))
	defer releaseVisited(visited)
	scratch := acquireHNSWScratch()
	defer releaseHNSWScratch(scratch)
	for layer := idx.maxLevel; layer > 0; layer-- {
		if err := checkCtx(ctx); err != nil {
			return nil, err
		}
		best := idx.searchLayer(query, queryNorm, current, 1, layer, cache, visited, scratch)
		if len(best) > 0 {
			current = best[0].rowIdx
		}
	}
	efSearch := chooseHNSWEfSearch(k)
	candidates := idx.searchLayer(query, queryNorm, current, efSearch, 0, cache, visited, scratch)
	resultHeap := &vecScoredHeap{}
	for _, sr := range candidates {
		pushTopK(resultHeap, sr.rowIdx, sr.distance, k)
	}
	if resultHeap.Len() < k {
		distFn := buildVecDistanceFunc(idx.metric, query, queryNorm, cache)
		return vecSearchTopK(ctx, idx.table.Rows, idx.dims, k, cache, distFn)
	}
	return topKFromHeap(resultHeap, k), nil
}

func (idx *vecHNSWIndex) searchLayer(query []float64, queryNorm float64, entry int, ef int, layer int, cache vecSearchColumnCacheEntry, visited []bool, scratch *vecHNSWScratch) []vecScoredRow {
	if entry < 0 || ef <= 0 || !idx.hasLayer(entry, layer) {
		return nil
	}
	dist, ok := rowDistance(idx.metric, query, queryNorm, cache, entry)
	if !ok {
		return nil
	}
	if len(visited) < len(cache.vectors) {
		visited = make([]bool, len(cache.vectors))
	}
	touched := scratch.touched[:0]
	markVisited := func(rowIdx int) bool {
		if rowIdx < 0 || rowIdx >= len(visited) || visited[rowIdx] {
			return false
		}
		visited[rowIdx] = true
		touched = append(touched, rowIdx)
		return true
	}
	markVisited(entry)
	defer func() {
		for _, rowIdx := range touched {
			visited[rowIdx] = false
		}
		scratch.touched = touched[:0]
	}()
	scratch.candidates = scratch.candidates[:0]
	scratch.results = scratch.results[:0]
	candidates := &scratch.candidates
	results := &scratch.results
	vecMinHeapPush(candidates, vecScoredRow{rowIdx: entry, distance: dist})
	pushTopK(results, entry, dist, ef)

	for candidates.Len() > 0 {
		nearest := vecMinHeapPop(candidates)
		if results.Len() >= ef && !vecScoredRowLess(nearest, (*results)[0]) {
			break
		}
		for _, nb := range idx.neighborLayer(nearest.rowIdx, layer) {
			if !markVisited(nb) {
				continue
			}
			dist, ok := rowDistance(idx.metric, query, queryNorm, cache, nb)
			if !ok {
				continue
			}
			if results.Len() < ef || vecScoredRowLess(vecScoredRow{rowIdx: nb, distance: dist}, (*results)[0]) {
				vecMinHeapPush(candidates, vecScoredRow{rowIdx: nb, distance: dist})
				pushTopK(results, nb, dist, ef)
			}
		}
	}
	return topKFromHeap(results, ef)
}

func (idx *vecHNSWIndex) addHNSWLink(a, b, layer int, cache vecSearchColumnCacheEntry) {
	if a == b || !idx.hasLayer(a, layer) || !idx.hasLayer(b, layer) {
		return
	}
	if !containsInt(idx.neighbors[a][layer], b) {
		idx.neighbors[a][layer] = append(idx.neighbors[a][layer], b)
		idx.pruneHNSWNeighbors(a, layer, cache)
	}
	if !containsInt(idx.neighbors[b][layer], a) {
		idx.neighbors[b][layer] = append(idx.neighbors[b][layer], a)
		idx.pruneHNSWNeighbors(b, layer, cache)
	}
}

func (idx *vecHNSWIndex) pruneHNSWNeighbors(rowIdx, layer int, cache vecSearchColumnCacheEntry) {
	nbs := idx.neighbors[rowIdx][layer]
	if len(nbs) <= vecHNSWM {
		return
	}
	query := cache.vectors[rowIdx]
	queryNorm := rowNormFor(idx.metric, cache, rowIdx)

	// addHNSWLink calls this on essentially every link insertion during the
	// sequential build loop, so this was the dominant cost of HNSW index
	// construction: a sort.Slice comparator recomputed rowDistance twice per
	// comparison (once for each side), on top of sort.Slice's own
	// reflection-based Swapper overhead. Precompute each neighbor's distance
	// once into a concrete vecScoredRow slice, then insertion-sort it in
	// place — nbs is bounded to vecHNSWM+1 elements here, so insertion sort
	// is both fast and avoids sort.Slice's overhead entirely.
	// vecScoredRowLess (distance ascending, ties broken by rowIdx ascending)
	// reproduces the exact ordering the old comparator computed, so which
	// neighbors survive pruning — and therefore search results — is
	// unchanged.
	scored := make([]vecScoredRow, len(nbs))
	for i, nb := range nbs {
		dist, ok := rowDistance(idx.metric, query, queryNorm, cache, nb)
		if !ok {
			dist = math.Inf(1)
		}
		scored[i] = vecScoredRow{rowIdx: nb, distance: dist}
	}
	for i := 1; i < len(scored); i++ {
		v := scored[i]
		j := i - 1
		for j >= 0 && vecScoredRowLess(v, scored[j]) {
			scored[j+1] = scored[j]
			j--
		}
		scored[j+1] = v
	}
	for i, sr := range scored {
		nbs[i] = sr.rowIdx
	}
	idx.neighbors[rowIdx][layer] = nbs[:vecHNSWM]
}

func (idx *vecHNSWIndex) hasLayer(rowIdx, layer int) bool {
	return rowIdx >= 0 && rowIdx < len(idx.neighbors) && layer >= 0 && layer < len(idx.neighbors[rowIdx])
}

func (idx *vecHNSWIndex) neighborLayer(rowIdx, layer int) []int {
	if !idx.hasLayer(rowIdx, layer) {
		return nil
	}
	return idx.neighbors[rowIdx][layer]
}

func validVectorRows(cache vecSearchColumnCacheEntry, dims int) []int {
	rows := make([]int, 0, len(cache.vectors))
	for i := range cache.vectors {
		if validCacheRow(cache, i, dims) {
			rows = append(rows, i)
		}
	}
	return rows
}

func validCacheRow(cache vecSearchColumnCacheEntry, rowIdx int, dims int) bool {
	return rowIdx >= 0 && rowIdx < len(cache.valid) && cache.valid[rowIdx] && len(cache.vectors[rowIdx]) == dims
}

func rowNorm(cache vecSearchColumnCacheEntry, rowIdx int) float64 {
	if rowIdx >= 0 && rowIdx < len(cache.norms) {
		return cache.norms[rowIdx]
	}
	return vectorL2Norm(cache.vectors[rowIdx])
}

// metricNeedsNorms reports whether the metric consumes vector norms.
// Only cosine does; computing norms for l2/manhattan/dot would double the
// per-row work in index build and search paths.
func metricNeedsNorms(metric string) bool {
	return metric == "cosine"
}

// rowNormFor returns the cached/computed norm only when the metric needs it.
func rowNormFor(metric string, cache vecSearchColumnCacheEntry, rowIdx int) float64 {
	if !metricNeedsNorms(metric) {
		return 0
	}
	return rowNorm(cache, rowIdx)
}

// rowDistance returns a ranking-only distance (see vectorRankingDistance):
// every call site is internal graph/list traversal, never a value exposed
// directly to a caller without first passing through
// vecSearchTopKWithIndex's finalize step.
func rowDistance(metric string, query []float64, queryNorm float64, cache vecSearchColumnCacheEntry, rowIdx int) (float64, bool) {
	if rowIdx < 0 || rowIdx >= len(cache.vectors) || !cache.valid[rowIdx] {
		return 0, false
	}
	return vectorRankingDistance(metric, cache.vectors[rowIdx], query, rowNormFor(metric, cache, rowIdx), queryNorm)
}

func centroidNorms(centroids [][]float64) []float64 {
	norms := make([]float64, len(centroids))
	for i, c := range centroids {
		norms[i] = vectorL2Norm(c)
	}
	return norms
}

// centroidNormsFor is centroidNorms gated by metricNeedsNorms: idx.centroidNorms
// is indexed unconditionally by nearestCentroid/search (idx.centroidNorms[i]),
// so it must always have one entry per centroid, but only cosine ever reads
// the values — l2/manhattan/dot ignore normA/normB entirely (vectorDistance).
// Returning a zero-filled slice of the same length for those metrics skips
// vectorL2Norm's sqrt per centroid per k-means iteration for no behavior
// change, mirroring rowNormFor's same gating for per-row norms.
func centroidNormsFor(metric string, centroids [][]float64) []float64 {
	if !metricNeedsNorms(metric) {
		return make([]float64, len(centroids))
	}
	return centroidNorms(centroids)
}

// nearestCentroid only ever returns an index, never the distance itself, so
// it uses the cheaper ranking-only distance (see vectorRankingDistance).
func nearestCentroid(metric string, vec []float64, vecNorm float64, centroids [][]float64, norms []float64) int {
	best := 0
	bestDist := math.MaxFloat64
	for i, c := range centroids {
		dist, ok := vectorRankingDistance(metric, c, vec, norms[i], vecNorm)
		if !ok {
			continue
		}
		if dist < bestDist {
			best = i
			bestDist = dist
		}
	}
	return best
}

func chooseIVFListCount(rows int) int {
	if rows <= 0 {
		return 0
	}
	nlist := int(math.Sqrt(float64(rows)) / 2)
	if nlist < 4 {
		nlist = min(rows, 4)
	}
	if nlist > 64 {
		nlist = 64
	}
	return nlist
}

func chooseIVFNProbe(nlist, k int) int {
	nprobe := int(math.Sqrt(float64(nlist))) + k/32
	if nprobe < 2 {
		nprobe = 2
	}
	if nprobe > 16 {
		nprobe = 16
	}
	if nprobe > nlist {
		nprobe = nlist
	}
	return nprobe
}

func hnswLevel(rowIdx int) int {
	x := uint64(rowIdx+1) * 11400714819323198485
	x ^= x >> 33
	x *= 0xff51afd7ed558ccd
	level := 0
	for level < vecHNSWMaxLevel && x&0x7 < 1 {
		level++
		x >>= 3
	}
	return level
}

// chooseHNSWEfSearch picks the graph-search beam width (ef) for a
// requested top-k. Below a floor of 16 it widens the beam so tiny k still
// gets a reasonable exploration budget; from 16 up it tracks k directly
// (ef == k), same as it always has for k in that range.
//
// It used to also clamp ef to a ceiling of 64, which meant any k > 64 got
// an ef strictly less than k. Since searchLayer's result set is bounded to
// at most ef entries (see searchLayer below), that made
// resultHeap.Len() < k in search() unconditionally true whenever k > 64,
// so every such call fell through to vecSearchTopK's full flat scan —
// paying for the graph traversal *and* the fallback, with the index
// providing no benefit at all for k > 64 (e.g. top-100/200 reranking
// queries). Dropping the ceiling and extending the existing ef == k
// relationship past 64 lets the graph search return k results directly
// for arbitrarily large k, exactly as it already did for k in [16, 64].
func chooseHNSWEfSearch(k int) int {
	ef := k
	if ef < 16 {
		ef = 16
	}
	return ef
}

func selectHNSWNeighbors(candidates []vecScoredRow, limit int) []vecScoredRow {
	if len(candidates) <= limit {
		return candidates
	}
	return candidates[:limit]
}

func containsInt(values []int, needle int) bool {
	for _, v := range values {
		if v == needle {
			return true
		}
	}
	return false
}
