package engine

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"testing"
	"time"

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

var vectorMathBenchmarkSink float64
var vectorMathBenchmarkIntSink int

func makeVectorMathBenchmarkInputs(dims int) ([]float64, []float64) {
	a := make([]float64, dims)
	b := make([]float64, dims)
	for i := range a {
		a[i] = math.Sin(float64(i)*0.11) * 0.75
		b[i] = math.Cos(float64(i)*0.07) * 0.5
	}
	return a, b
}

func BenchmarkVectorDot768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorDot(a, vecB)
	}
}

func BenchmarkVectorDotUnrolled768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorDotUnrolled(a, vecB)
	}
}

func BenchmarkVectorL2Squared768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorL2Squared(a, vecB)
	}
}

func BenchmarkVectorL2SquaredUnrolled768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorL2SquaredUnrolled(a, vecB)
	}
}

func BenchmarkVectorL1Distance768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorL1Distance(a, vecB)
	}
}

func BenchmarkVectorL1DistanceUnrolled768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkSink = vectorL1Unrolled(a, vecB)
	}
}

// BenchmarkVectorHammingDistance768 measures vectorHammingDistance (which
// dispatches to the portable 4-way-unrolled vectorHammingUnrolled — see the
// design-rationale comment on vectorHammingDistance in vector_math.go for why
// this metric intentionally has no hand-rolled SIMD kernel) against a plain,
// non-unrolled reference loop (naiveHamming, already defined in
// vector_math_unroll_test.go for correctness cross-checking and reused here).
// Run side by side so any future change to either implementation shows up as
// a relative regression/improvement, not just an absolute number.
func BenchmarkVectorHammingDistance768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkIntSink = vectorHammingDistance(a, vecB)
	}
}

func BenchmarkVectorHammingNaive768(b *testing.B) {
	a, vecB := makeVectorMathBenchmarkInputs(768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		vectorMathBenchmarkIntSink = naiveHamming(a, vecB)
	}
}

// makeCentroidBenchmarkInputs builds `vectors` deterministic 768-ish-dim
// vectors for VEC_CENTROID-style accumulation benchmarks below.
func makeCentroidBenchmarkInputs(vectors, dims int) [][]float64 {
	vecs := make([][]float64, vectors)
	for v := 0; v < vectors; v++ {
		vec := make([]float64, dims)
		for i := range vec {
			vec[i] = math.Sin(float64(i)*0.11+float64(v)) * 0.75
		}
		vecs[v] = vec
	}
	return vecs
}

// BenchmarkVectorCentroidAccumulate768x8 and
// BenchmarkVectorCentroidAccumulateNaive768x8 mirror evalVecCentroid's own
// usage of vectorAccumulateUnrolled: a fresh zeroed accumulator per call,
// with each of 8 768-dim source vectors folded into it in turn (see
// VEC_CENTROID in vector_functions.go). naiveAccumulate is the trivial
// reference loop already defined in vector_math_unroll_test.go for
// correctness cross-checking, reused here for the timing comparison.
func BenchmarkVectorCentroidAccumulate768x8(b *testing.B) {
	vecs := makeCentroidBenchmarkInputs(8, 768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		out := make([]float64, 768)
		for _, v := range vecs {
			vectorAccumulateUnrolled(out, v)
		}
		vectorMathBenchmarkSink = out[0]
	}
}

func BenchmarkVectorCentroidAccumulateNaive768x8(b *testing.B) {
	vecs := makeCentroidBenchmarkInputs(8, 768)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		out := make([]float64, 768)
		for _, v := range vecs {
			naiveAccumulate(out, v)
		}
		vectorMathBenchmarkSink = out[0]
	}
}

// BenchmarkVectorDotKernelBySize guards the amd64 kernel dispatch decision in
// vectorDotKernel (vector_math_amd64.go): that function used to fall back to
// the portable vectorDotUnrolled loop for vectors shorter than 128 elements,
// on the assumption that SSE2 setup overhead wasn't worth it below that size.
// Measuring across the dimension sizes real embedding models actually use
// (16 up to 768) showed the SSE2 kernel winning at every single size,
// including 16 — the assembly kernel's own tail loop already handles
// short/odd-length inputs, so there was no setup cost being amortized and
// the threshold only ever left performance on the table. Run both variants
// side by side here so a future change to that threshold (or a new
// portable-loop "improvement") shows up as a regression instead of silently
// reintroducing the bad cutover.
func BenchmarkVectorDotKernelBySize(b *testing.B) {
	for _, dims := range []int{16, 32, 64, 128, 256, 768} {
		a, vecB := makeVectorMathBenchmarkInputs(dims)
		b.Run(fmt.Sprintf("unrolled/%d", dims), func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				vectorMathBenchmarkSink = vectorDotUnrolled(a, vecB)
			}
		})
		b.Run(fmt.Sprintf("kernel/%d", dims), func(b *testing.B) {
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				vectorMathBenchmarkSink = vectorDotKernel(a, vecB)
			}
		})
	}
}

func makeVecSearchBenchmarkTable(rows, dims int) *storage.DB {
	db := storage.NewDB()
	table := storage.NewTable("vec_docs", []storage.Column{
		{Name: "id", Type: storage.IntType},
		{Name: "embedding", Type: storage.VectorType},
	}, false)

	for i := 0; i < rows; i++ {
		vec := make([]float64, dims)
		for d := 0; d < dims; d++ {
			angle := float64(i*d) / float64(rows*dims+1)
			vec[d] = math.Sin(angle) * 0.5
		}
		table.Rows = append(table.Rows, []any{i, vec})
	}

	if err := db.Put("default", table); err != nil {
		panic(err)
	}

	return db
}

func benchmarkVecSearch(b *testing.B, metric string) {
	db := makeVecSearchBenchmarkTable(4096, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = float64(i) / 64.0
	}

	args := []Expr{
		&Literal{Val: "vec_docs"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: 5},
		&Literal{Val: metric},
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := fn.Execute(context.Background(), args, env, Row{})
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 5 {
			b.Fatalf("expected top-5 results, got %d", len(rs.Rows))
		}
	}
}

func benchmarkVecSearchIndexed(b *testing.B, indexMode string) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	args := []Expr{
		&Literal{Val: "rag_hybrid"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: 20},
		&Literal{Val: "cosine"},
		&Literal{Val: indexMode},
	}

	if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
		b.Fatal(err)
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := fn.Execute(context.Background(), args, env, Row{})
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) == 0 || len(rs.Rows) > 20 {
			b.Fatalf("expected up to 20 results, got %d", len(rs.Rows))
		}
	}
}

// benchmarkVecSearchIndexedK mirrors benchmarkVecSearchIndexed but takes an
// explicit k, so callers can measure VEC_SEARCH well above the 64-result
// ceiling chooseHNSWEfSearch used to hard-cap its graph-search beam (ef) at
// — see BenchmarkVecSearchCosineTopK100_HNSWCached/
// BenchmarkVecSearchCosineTopK200_HNSWCached below.
func benchmarkVecSearchIndexedK(b *testing.B, indexMode string, k int) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	args := []Expr{
		&Literal{Val: "rag_hybrid"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: k},
		&Literal{Val: "cosine"},
		&Literal{Val: indexMode},
	}

	if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
		b.Fatal(err)
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := fn.Execute(context.Background(), args, env, Row{})
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) == 0 || len(rs.Rows) > k {
			b.Fatalf("expected up to %d results, got %d", k, len(rs.Rows))
		}
	}
}

func makeWhereAndBenchmarkTable(rows, dims int) *storage.DB {
	db := storage.NewDB()
	table := storage.NewTable("rag_docs", []storage.Column{
		{Name: "id", Type: storage.IntType},
		{Name: "score", Type: storage.IntType},
		{Name: "embedding", Type: storage.VectorType},
	}, false)

	for i := 0; i < rows; i++ {
		vec := make([]float64, dims)
		for d := 0; d < dims; d++ {
			angle := 0.08*float64(i) + 0.17*float64(d)
			vec[d] = math.Sin(angle)
		}
		table.Rows = append(table.Rows, []any{i, i, vec})
	}

	if err := db.Put("default", table); err != nil {
		panic(err)
	}

	return db
}

func makeRAGHybridBenchmarkTable(rows, dims int) *storage.DB {
	db := storage.NewDB()
	table := storage.NewTable("rag_hybrid", []storage.Column{
		{Name: "id", Type: storage.IntType},
		{Name: "score", Type: storage.IntType},
		{Name: "quality", Type: storage.FloatType},
		{Name: "created_at", Type: storage.TextType},
		{Name: "embedding", Type: storage.VectorType},
	}, false)

	base := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)
	for i := 0; i < rows; i++ {
		vec := make([]float64, dims)
		for d := 0; d < dims; d++ {
			angle := 0.08*float64(i) + 0.17*float64(d)
			vec[d] = math.Sin(angle)
		}
		ts := base.Add(-time.Duration(i) * time.Minute).Format("2006-01-02 15:04:05")
		quality := 0.5 + 0.5*math.Sin(float64(i)*0.03)
		table.Rows = append(table.Rows, []any{i, i, quality, ts, vec})
	}

	if err := db.Put("default", table); err != nil {
		panic(err)
	}
	return db
}

func makeRAGChunkBenchmarkTable(rows, dims int) *storage.DB {
	db := storage.NewDB()
	table := storage.NewTable("rag_chunks", []storage.Column{
		{Name: "doc_id", Type: storage.TextType},
		{Name: "chunk_index", Type: storage.IntType},
		{Name: "chunk_text", Type: storage.TextType},
		{Name: "quality", Type: storage.FloatType},
		{Name: "created_at", Type: storage.TextType},
		{Name: "embedding", Type: storage.VectorType},
	}, false)

	base := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC)
	for i := 0; i < rows; i++ {
		vec := make([]float64, dims)
		for d := 0; d < dims; d++ {
			angle := 0.08*float64(i) + 0.17*float64(d)
			vec[d] = math.Sin(angle)
		}
		docID := "doc-" + strconv.Itoa(i/12)
		chunkIndex := i % 12
		chunkText := "chunk " + strconv.Itoa(i)
		quality := 0.5 + 0.5*math.Sin(float64(i)*0.03)
		ts := base.Add(-time.Duration(i) * time.Minute).Format("2006-01-02 15:04:05")
		table.Rows = append(table.Rows, []any{docID, chunkIndex, chunkText, quality, ts, vec})
	}

	if err := db.Put("default", table); err != nil {
		panic(err)
	}
	return db
}

func benchmarkWhereAndVector(b *testing.B, vectorFirst bool) {
	db := makeWhereAndBenchmarkTable(12000, 32)
	query := make([]float64, 32)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}

	var stmt Statement
	if vectorFirst {
		stmt = mustParse(`
			SELECT id, score
			FROM rag_docs
			WHERE VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) > -1.0
				AND score > 9000
		`)
	} else {
		stmt = mustParse(`
			SELECT id, score
			FROM rag_docs
			WHERE score > 9000
				AND VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) > -1.0
		`)
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 2999 {
			b.Fatalf("expected 2999 rows, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkVecSearchCosineTopK(b *testing.B) {
	benchmarkVecSearch(b, "cosine")
}

func BenchmarkVecSearchL2TopK(b *testing.B) {
	benchmarkVecSearch(b, "l2")
}

func BenchmarkVecSearchCosineTopK_IVFCached(b *testing.B) {
	benchmarkVecSearchIndexed(b, "ivf")
}

func BenchmarkVecSearchCosineTopK_HNSWCached(b *testing.B) {
	benchmarkVecSearchIndexed(b, "hnsw")
}

// BenchmarkVecSearchCosineTopK100_HNSWCached and
// BenchmarkVecSearchCosineTopK200_HNSWCached exercise VEC_SEARCH's HNSW path
// at k values above chooseHNSWEfSearch's old fixed ceiling of 64. Before
// that ceiling was removed, vecHNSWIndex.search's graph-search beam
// (efSearch) topped out at 64 regardless of k, so any k > 64 always
// produced fewer than k graph candidates and fell through to
// vecSearchTopK's full flat scan on every single call — paying for both
// the graph traversal and the linear scan, with the index providing zero
// benefit for exactly this "top-100/200 for downstream reranking" RAG
// pattern. Their _IVFCached siblings are measured the same way as a
// same-corpus reference point (IVF was never affected by this bug).
func BenchmarkVecSearchCosineTopK100_HNSWCached(b *testing.B) {
	benchmarkVecSearchIndexedK(b, "hnsw", 100)
}

func BenchmarkVecSearchCosineTopK100_IVFCached(b *testing.B) {
	benchmarkVecSearchIndexedK(b, "ivf", 100)
}

func BenchmarkVecSearchCosineTopK200_HNSWCached(b *testing.B) {
	benchmarkVecSearchIndexedK(b, "hnsw", 200)
}

func BenchmarkVecSearchCosineTopK200_IVFCached(b *testing.B) {
	benchmarkVecSearchIndexedK(b, "ivf", 200)
}

func BenchmarkVecSearchIndexModesSameTable(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}

	for _, indexMode := range []string{"flat", "ivf", "hnsw"} {
		indexMode := indexMode
		b.Run(indexMode, func(b *testing.B) {
			args := []Expr{
				&Literal{Val: "rag_hybrid"},
				&Literal{Val: "embedding"},
				&Literal{Val: query},
				&Literal{Val: 20},
				&Literal{Val: "cosine"},
				&Literal{Val: indexMode},
			}
			if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
				b.Fatal(err)
			}

			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				rs, err := fn.Execute(context.Background(), args, env, Row{})
				if err != nil {
					b.Fatal(err)
				}
				if len(rs.Rows) == 0 || len(rs.Rows) > 20 {
					b.Fatalf("expected up to 20 results, got %d", len(rs.Rows))
				}
			}
		})
	}
}

// BenchmarkVecSearchIndexModesSameTableL2 is BenchmarkVecSearchIndexModesSameTable's
// l2-metric twin: l2's flat/IVF/HNSW paths all defer math.Sqrt to the final
// (at most k) results instead of paying it per candidate/edge (see
// vectorRankingDistance/finalizeVecScoredRows in vector_math.go/
// vector_index.go), so this demonstrates that win across all three index
// modes the cosine-only version above can't exercise.
func BenchmarkVecSearchIndexModesSameTableL2(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}

	for _, indexMode := range []string{"flat", "ivf", "hnsw"} {
		indexMode := indexMode
		b.Run(indexMode, func(b *testing.B) {
			args := []Expr{
				&Literal{Val: "rag_hybrid"},
				&Literal{Val: "embedding"},
				&Literal{Val: query},
				&Literal{Val: 20},
				&Literal{Val: "l2"},
				&Literal{Val: indexMode},
			}
			if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
				b.Fatal(err)
			}

			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				rs, err := fn.Execute(context.Background(), args, env, Row{})
				if err != nil {
					b.Fatal(err)
				}
				if len(rs.Rows) == 0 || len(rs.Rows) > 20 {
					b.Fatalf("expected up to 20 results, got %d", len(rs.Rows))
				}
			}
		})
	}
}

func BenchmarkWhereVectorAndSimpleCondition_VectorThenScalar(b *testing.B) {
	benchmarkWhereAndVector(b, true)
}

func BenchmarkWhereVectorAndSimpleCondition_SimpleThenVector(b *testing.B) {
	benchmarkWhereAndVector(b, false)
}

func BenchmarkHybridOrderByVectorAndRecency(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		SELECT id,
		       RAG_HYBRID_SCORE(
		           VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')),
		           created_at,
		           30,
		           0.7,
		           '2026-01-01 01:00:00'
		       ) AS rag_score
		FROM rag_hybrid
		WHERE VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) > -1.0
		  AND RECENCY_SCORE(created_at, 30, '2026-01-01 01:00:00') > 0.0
		ORDER BY rag_score DESC
		LIMIT 20
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 20 {
			b.Fatalf("expected 20 results, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkRAGRankScoreOrderByLimit(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		SELECT id,
		       RAG_RANK_SCORE(
		           VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')),
		           created_at,
		           30,
		           quality,
		           0.65,
		           0.25,
		           0.10,
		           '2026-01-01 01:00:00'
		       ) AS rag_score
		FROM rag_hybrid
		ORDER BY rag_score DESC
		LIMIT 20
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 20 {
			b.Fatalf("expected 20 results, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkRAGContextFromTopK(b *testing.B) {
	db := makeRAGChunkBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		WITH topk AS (
			SELECT doc_id, chunk_index, _vec_rank
			FROM VEC_SEARCH('rag_chunks', 'embedding', VEC_FROM_JSON('` + string(queryJSON) + `'), 20, 'cosine')
		)
		SELECT doc_id, chunk_index, chunk_text, _hit_rank, _context_offset
		FROM RAG_CONTEXT_FROM('rag_chunks', 'doc_id', 'chunk_index', 'topk', 'doc_id', 'chunk_index', 1, 1)
		ORDER BY _context_rank
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) == 0 {
			b.Fatal("expected context rows")
		}
	}
}

func BenchmarkRAGContextSingle(b *testing.B) {
	db := makeRAGChunkBenchmarkTable(12000, 64)
	stmt := mustParse(`
		SELECT doc_id, chunk_index, chunk_text, _context_offset
		FROM RAG_CONTEXT('rag_chunks', 'doc_id', 'chunk_index', 'doc-500', 6, 2, 2)
		ORDER BY _context_rank
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 5 {
			b.Fatalf("expected 5 context rows, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkOrderByVectorLimit_NoWhere(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		SELECT VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) AS sim
		FROM rag_hybrid
		ORDER BY sim DESC
		LIMIT 20
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 20 {
			b.Fatalf("expected 20 results, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkOrderByVectorLimit_WithWhereAndScalar(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		SELECT VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) AS sim
		FROM rag_hybrid
		WHERE score > 9000
		ORDER BY sim DESC
		LIMIT 20
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 20 {
			b.Fatalf("expected 20 results, got %d", len(rs.Rows))
		}
	}
}

func BenchmarkCompareTopK_VecSearchVsOrderBy(b *testing.B) {
	db := makeRAGHybridBenchmarkTable(12000, 64)
	query := make([]float64, 64)
	for i := range query {
		query[i] = math.Cos(0.08*float64(i) + 0.5)
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	queryJSONStr := string(queryJSON)

	stmtOrder := mustParse(`
		SELECT VEC_COSINE_SIMILARITY(embedding, VEC_FROM_JSON('` + queryJSONStr + `')) AS sim
		FROM rag_hybrid
		ORDER BY sim DESC
		LIMIT 20
	`)
	stmtVec := mustParse(`
		SELECT * FROM VEC_SEARCH('rag_hybrid', 'embedding', VEC_FROM_JSON('` + queryJSONStr + `'), 20, 'cosine')
	`)

	b.Run("order_by_vector", func(b *testing.B) {
		b.ReportAllocs()
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			rs, err := Execute(context.Background(), db, "default", stmtOrder)
			if err != nil {
				b.Fatal(err)
			}
			if len(rs.Rows) != 20 {
				b.Fatalf("expected 20 results, got %d", len(rs.Rows))
			}
		}
	})

	b.Run("vec_search", func(b *testing.B) {
		b.ReportAllocs()
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			rs, err := Execute(context.Background(), db, "default", stmtVec)
			if err != nil {
				b.Fatal(err)
			}
			if len(rs.Rows) != 20 {
				b.Fatalf("expected 20 results, got %d", len(rs.Rows))
			}
		}
	})
}

// BenchmarkVecSearchConcurrent measures VEC_SEARCH under concurrent callers
// with the result cache and analytics both at their documented default (off)
// — the state every embedding application starts in. Rows are kept below
// vecSearchParallelMinRows (4096) so each call takes the single-goroutine
// path inside vecSearchTopK; that isolates vecQueryCacheState's (former)
// global-mutex overhead from the per-query worker-pool fan-out, which is a
// separate concern already covered by the other VEC_SEARCH benchmarks. A
// small table also keeps each call's own useful work short, so two
// process-wide Lock/Unlock round trips per call are a large fraction of wall
// time and any contention between goroutines shows up clearly rather than
// getting lost in scan noise. Run with -cpu=1,2,4,8 (or more) to see whether
// throughput scales with cores, which it cannot while every goroutine
// serializes on one shared lock.
func BenchmarkVecSearchConcurrent(b *testing.B) {
	db := makeVecSearchBenchmarkTable(512, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = float64(i) / 64.0
	}
	args := []Expr{
		&Literal{Val: "vec_docs"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: 5},
		&Literal{Val: "cosine"},
	}
	ConfigureVectorCache(VectorCacheConfig{})
	// Warm the column cache before timing, same reasoning as the flat-scan
	// benchmarks below: without this, whichever calibration attempt ends up
	// reported pays one unpredictable cold-build cost inside RunParallel,
	// where it also races several goroutines against getVecColumnCache's
	// build-coalescing channel instead of the mutex this benchmark targets.
	if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
		b.Fatal(err)
	}

	b.ReportAllocs()
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			rs, err := fn.Execute(context.Background(), args, env, Row{})
			if err != nil {
				b.Fatal(err)
			}
			if len(rs.Rows) != 5 {
				b.Fatal("expected top-5 results")
			}
		}
	})
}

// BenchmarkVecSearchFlatScanLargeEmbedding exercises the flat-scan column
// cache at a corpus size and embedding width representative of a real RAG
// deployment (50k chunks, 768 dims — a common BERT/sentence-transformer
// size), large enough that per-row cache/TLB misses from a scattered
// [][]float64 column cache dominate over the small-table benchmarks above
// (4096/12000 rows at 64 dims, whose ~512B rows can still fit close together
// under the allocator regardless of layout).
func BenchmarkVecSearchFlatScanLargeEmbedding(b *testing.B) {
	db := makeVecSearchBenchmarkTable(50000, 768)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 768)
	for i := range query {
		query[i] = float64(i) / 768.0
	}
	args := []Expr{
		&Literal{Val: "vec_docs"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: 10},
		&Literal{Val: "cosine"},
	}

	// Warm the column cache before timing starts (matches benchmarkVecSearchIndexed's
	// pattern above): the one-time cold build of a 50k-row/768-dim column is a
	// separate cost from steady-state scan latency, and go test's benchtime-based
	// calibration re-invokes this whole function (rebuilding db and the cache)
	// on every calibration attempt — without a warm-up call, whichever attempt
	// happens to become the reported one mixes in one cold build at an
	// unpredictable weight, making ns/op and allocs/op incomparable across runs.
	if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
		b.Fatal(err)
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := fn.Execute(context.Background(), args, env, Row{})
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 10 {
			b.Fatalf("expected top-10 results, got %d", len(rs.Rows))
		}
	}
}

// BenchmarkVecSearchFlatScanManySmallEmbeddings targets the opposite corner
// of the memory-layout tradeoff from BenchmarkVecSearchFlatScanLargeEmbedding:
// many rows (200k) with a small embedding (64 dims, 512 bytes/row). The full
// column here (~100MB) is far larger than any CPU cache, and each row's own
// payload is small, so the fixed per-row cost of following one more pointer
// indirection to find a scattered allocation is a much larger fraction of
// that row's total cost than it is for a 768-dim/6KB row — this is where a
// contiguous backing buffer should show its clearest win.
func BenchmarkVecSearchFlatScanManySmallEmbeddings(b *testing.B) {
	db := makeVecSearchBenchmarkTable(200000, 64)
	fn := &VecSearchTableFunc{}
	env := ExecEnv{ctx: context.Background(), tenant: "default", db: db}
	query := make([]float64, 64)
	for i := range query {
		query[i] = float64(i) / 64.0
	}
	args := []Expr{
		&Literal{Val: "vec_docs"},
		&Literal{Val: "embedding"},
		&Literal{Val: query},
		&Literal{Val: 10},
		&Literal{Val: "cosine"},
	}
	if _, err := fn.Execute(context.Background(), args, env, Row{}); err != nil {
		b.Fatal(err)
	}

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := fn.Execute(context.Background(), args, env, Row{})
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 10 {
			b.Fatalf("expected top-10 results, got %d", len(rs.Rows))
		}
	}
}

// makeHammingBenchmarkTable builds a table of binary-quantized-style
// embeddings (each element is exactly 0.0 or 1.0, as VEC_BINARY_QUANTIZE
// would produce) so VEC_HAMMING_DISTANCE comparisons below are meaningful
// rather than degenerating to "everything differs".
func makeHammingBenchmarkTable(rows, dims int) *storage.DB {
	db := storage.NewDB()
	table := storage.NewTable("hamming_docs", []storage.Column{
		{Name: "id", Type: storage.IntType},
		{Name: "embedding", Type: storage.VectorType},
	}, false)

	for i := 0; i < rows; i++ {
		vec := make([]float64, dims)
		for d := 0; d < dims; d++ {
			// Deterministic pseudo-random bit pattern per row/dim.
			if math.Sin(float64(i)*0.31+float64(d)*0.53) > 0 {
				vec[d] = 1.0
			}
		}
		table.Rows = append(table.Rows, []any{i, vec})
	}

	if err := db.Put("default", table); err != nil {
		panic(err)
	}
	return db
}

// BenchmarkOrderByHammingDistanceLimit measures the one usage pattern that
// actually puts VEC_HAMMING_DISTANCE (and, transitively, the unrolled
// portable kernel it dispatches to) on a per-row hot path today: a full
// table scan computing the distance for every row via ORDER BY ... LIMIT,
// mirroring BenchmarkOrderByVectorLimit_NoWhere above but for Hamming
// distance over binary-quantized embeddings instead of cosine similarity.
// This is the benchmark that's actually representative of real-world
// relevance, since none of VEC_HAMMING_DISTANCE/VEC_CENTROID are confirmed
// to sit on any other hot path in tinySQL today.
func BenchmarkOrderByHammingDistanceLimit(b *testing.B) {
	const rows, dims = 12000, 768
	db := makeHammingBenchmarkTable(rows, dims)
	query := make([]float64, dims)
	for i := range query {
		if math.Cos(float64(i)*0.17) > 0 {
			query[i] = 1.0
		}
	}
	queryJSON, err := json.Marshal(query)
	if err != nil {
		b.Fatal(err)
	}
	stmt := mustParse(`
		SELECT id, VEC_HAMMING_DISTANCE(embedding, VEC_FROM_JSON('` + string(queryJSON) + `')) AS hdist
		FROM hamming_docs
		ORDER BY hdist ASC
		LIMIT 20
	`)

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rs, err := Execute(context.Background(), db, "default", stmt)
		if err != nil {
			b.Fatal(err)
		}
		if len(rs.Rows) != 20 {
			b.Fatalf("expected 20 results, got %d", len(rs.Rows))
		}
	}
}
