// Command bench measures the index, so §7's numbers are recorded
// measurements rather than assertions.
//
// A review of draft 2 pointed out that every figure in §7 — 3.2 GFLOP/s,
// 104 MB, 19 s — had no benchmark, command or machine behind it, and that
// the FTS size could not be reconciled with the stated schema, because a
// default FTS5 table keeps its own copy of the inserted text. Both are
// fair, and both are answerable by measuring instead of arguing.
package main

import (
	"crypto/rand"
	"database/sql"
	"flag"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	_ "modernc.org/sqlite"

	"amadan.net/oficina/memoria/internal/memoria/chunk"
)

func main() {
	corpus := flag.String("corpus", "", "directory of real text to build the corpus from")
	target := flag.Int("chunks", 90000, "how many chunks to grow the corpus to")
	flag.Parse()

	fmt.Printf("machine: %s/%s, %d cpu, go %s\n\n", runtime.GOOS, runtime.GOARCH, runtime.NumCPU(), runtime.Version())

	passages := realPassages(*corpus)
	if len(passages) == 0 {
		fmt.Fprintln(os.Stderr, "no corpus found; pass --corpus")
		os.Exit(1)
	}
	avgWords := 0
	avgBytes := 0
	for _, p := range passages {
		avgWords += len(strings.Fields(p))
		avgBytes += len(p)
	}
	fmt.Printf("real corpus: %d passages, mean %d words / %d bytes per passage\n",
		len(passages), avgWords/len(passages), avgBytes/len(passages))

	// Grow to the target by cycling the real passages. Cycling repeats
	// vocabulary, which makes the postings list SMALLER than genuinely
	// varied English would — so the FTS figure below is optimistic, and
	// says so.
	grown := make([]string, 0, *target)
	for i := 0; i < *target; i++ {
		grown = append(grown, passages[i%len(passages)])
	}
	totalBytes := 0
	for _, p := range grown {
		totalBytes += len(p)
	}
	fmt.Printf("grown corpus: %d chunks, %.1f MB of text\n\n", len(grown), float64(totalBytes)/(1<<20))

	for _, contentless := range []bool{false, true} {
		measureFTS(grown, contentless)
	}
	measureCosine()
}

func measureFTS(chunks []string, contentless bool) {
	label := "default (keeps its own copy of the text)"
	schema := `CREATE VIRTUAL TABLE c USING fts5(body, tokenize='porter unicode61');`
	if contentless {
		label = "contentless (content='')"
		schema = `CREATE VIRTUAL TABLE c USING fts5(body, content='', tokenize='porter unicode61');`
	}
	dir, err := os.MkdirTemp("/home/paulca/memoriadrive", "bench")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(dir)
	path := filepath.Join(dir, "fts.db")
	db, err := sql.Open("sqlite", "file:"+path)
	if err != nil {
		panic(err)
	}
	defer db.Close()
	if _, err := db.Exec(schema); err != nil {
		panic(err)
	}

	start := time.Now()
	tx, _ := db.Begin()
	st, err := tx.Prepare(`INSERT INTO c (body) VALUES (?)`)
	if err != nil {
		panic(err)
	}
	for _, ch := range chunks {
		if _, err := st.Exec(ch); err != nil {
			panic(err)
		}
	}
	st.Close()
	if err := tx.Commit(); err != nil {
		panic(err)
	}
	if _, err := db.Exec(`INSERT INTO c(c) VALUES('optimize');`); err != nil {
		panic(err)
	}
	build := time.Since(start)
	db.Exec(`PRAGMA wal_checkpoint(TRUNCATE);`)

	var size int64
	for _, suffix := range []string{"", "-wal"} {
		if fi, err := os.Stat(path + suffix); err == nil {
			size += fi.Size()
		}
	}

	// A query, warm.
	qstart := time.Now()
	var n int
	row := db.QueryRow(`SELECT count(*) FROM c WHERE c MATCH '"revenue" OR "sealed"'`)
	if err := row.Scan(&n); err != nil {
		fmt.Printf("  query failed: %v\n", err)
	}
	q := time.Since(qstart)

	fmt.Printf("FTS5 %s\n", label)
	fmt.Printf("  index on disk : %.1f MB (%d bytes/chunk)\n", float64(size)/(1<<20), size/int64(len(chunks)))
	fmt.Printf("  build time    : %s\n", build.Round(time.Millisecond))
	fmt.Printf("  one BM25 query: %s (%d rows matched)\n\n", q.Round(time.Microsecond), n)
}

// measureCosine times a brute-force pass, the number §7's latency table
// is derived from.
func measureCosine() {
	fmt.Println("brute-force cosine, one goroutine, pure Go, no SIMD")
	for _, c := range []struct{ n, d int }{{10000, 512}, {90000, 512}, {100000, 768}, {250000, 512}} {
		vecs := make([]float32, c.n*c.d)
		buf := make([]byte, len(vecs)*4)
		rand.Read(buf)
		for i := range vecs {
			vecs[i] = float32(int8(buf[i*4])) / 128
		}
		q := make([]float32, c.d)
		for i := range q {
			q[i] = float32(i%17) / 17
		}
		// Three runs, report the best — a slow run measures the
		// scheduler, not the arithmetic.
		best := time.Duration(math.MaxInt64)
		for run := 0; run < 3; run++ {
			start := time.Now()
			var bestScore float32
			for i := 0; i+c.d <= len(vecs); i += c.d {
				var dot float32
				row := vecs[i : i+c.d]
				for j, v := range row {
					dot += v * q[j]
				}
				if dot > bestScore {
					bestScore = dot
				}
			}
			if e := time.Since(start); e < best {
				best = e
			}
		}
		flops := float64(2 * c.n * c.d)
		fmt.Printf("  %7d chunks x %d dims: %8s  (%.1f GFLOP/s)  vectors %.0f MB f32 / %.0f MB int8\n",
			c.n, c.d, best.Round(time.Microsecond), flops/best.Seconds()/1e9,
			float64(c.n*c.d*4)/(1<<20), float64(c.n*c.d)/(1<<20))
	}
}

// realPassages reads .md files and chunks them with MEMORIA'S OWN
// chunker, so the measurement describes the real product rather than an
// approximation of it. An earlier version split on blank lines and
// reported a mean of 81 words per passage, which is a paragraph, not a
// memoria chunk.
func realPassages(dir string) []string {
	var out []string
	if dir == "" {
		return nil
	}
	filepath.Walk(dir, func(p string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return nil
		}
		if !strings.HasSuffix(p, ".md") || strings.Contains(p, "/.") || strings.Contains(p, "node_modules") {
			return nil
		}
		raw, err := os.ReadFile(p)
		if err != nil {
			return nil
		}
		for _, c := range chunk.Split(string(raw)) {
			out = append(out, c.Indexable())
		}
		return nil
	})
	return out
}