package store

import (
	"context"
	"fmt"
	"math"
	"math/rand"
	"path/filepath"
	"testing"
)

// benchDim is the resident embedder's width (multilingual-e5-small, 384), so
// the per-row decode cost the benchmark measures is the real one.
const benchDim = 384

// seedMemVectors fills a fresh store with n L2-normalized rows carrying a meta
// blob the size recall actually stores — the note text plus its type — because
// the cost this benchmark exists to measure is unmarshalling that blob for
// every row when only topK survivors need it.
func seedMemVectors(tb testing.TB, n int) *MemoryStore {
	tb.Helper()
	path := filepath.Join(tb.TempDir(), "mem_bench.db")
	st, err := Open(context.Background(), path)
	if err != nil {
		tb.Fatalf("Open: %v", err)
	}
	tb.Cleanup(func() { _ = st.Close() })
	m := st.VectorMemory()

	rng := rand.New(rand.NewSource(1))
	ctx := context.Background()
	for i := 0; i < n; i++ {
		if err := m.Insert(ctx, fmt.Sprintf("note:%d", i), randUnitVec(rng, benchDim), map[string]string{
			"type": "note",
			"text": fmt.Sprintf("заметка номер %d о том, что надо не забыть сделать на неделе", i),
		}); err != nil {
			tb.Fatalf("Insert %d: %v", i, err)
		}
	}
	return m
}

func randUnitVec(rng *rand.Rand, dim int) []float32 {
	v := make([]float32, dim)
	var norm float64
	for i := range v {
		f := rng.NormFloat64()
		v[i] = float32(f)
		norm += f * f
	}
	norm = math.Sqrt(norm)
	for i := range v {
		v[i] = float32(float64(v[i]) / norm)
	}
	return v
}

// BenchmarkMemoryStoreSearch measures one recall query against a store of n
// rows. Row counts bracket the documented scale: 1000 is a plausible today,
// 10000 is the "thousands, not millions" ceiling the type doc claims a full
// scan is fine at.
func BenchmarkMemoryStoreSearch(b *testing.B) {
	for _, n := range []int{1000, 10000} {
		b.Run(fmt.Sprintf("rows=%d", n), func(b *testing.B) {
			m := seedMemVectors(b, n)
			q := randUnitVec(rand.New(rand.NewSource(2)), benchDim)
			ctx := context.Background()
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				if _, err := m.Search(ctx, q, 10); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}