// rag.go adds retrieval-augmented search on top of the plain
// read-a-whole-file skill lookup in skillkit.go. The read_office_skill
// tool (skillkit.Read) is the right fit when the model already knows
// which format it needs and wants that skill's full guidance before a
// multi-step task — but for a narrow, specific question ("can a cell
// hold a formula", "how do I bold part of a paragraph"), returning an
// entire 3-4KB skill file just to answer it is the same imprecision
// Claude Code's own skills avoid by splitting a large skill into a
// short SKILL.md plus separate reference files loaded on demand.
// Because every skill here already follows one convention — a "##"
// heading per topic — each section is naturally a single, coherent
// chunk, and because the whole corpus is a few dozen short sections
// fixed at compile time (go:embed, never changes between deploys),
// embedding it once per process and ranking in memory is simpler and
// faster than standing up a vector-DB round trip for content this
// small and this static. search_codebase, by contrast, indexes a
// workspace's actual source files — large, per-user, and changing —
// which is exactly the shape services/knowledge's Postgres/pgvector
// pipeline exists for; that pipeline would be the wrong tool here, not
// a missing piece.
package skillkit

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strings"
	"sync"
)

// Embedder is the subset of *vllmkit.Client this package needs —
// declared as a narrow interface (rather than importing vllmkit
// directly) so the caller supplies whatever embedding client it
// already has, and so this package stays testable without a live
// vLLM instance.
type Embedder interface {
	Embed(ctx context.Context, text string) ([]float32, error)
}

// Chunk is one semantically coherent section of a skill file — in
// practice, one "## Heading" and the text under it (or the file's
// title/intro, for the section before the first heading).
type Chunk struct {
	Skill   string // the skill this section belongs to, e.g. "docx"
	Heading string // the section's heading, e.g. "Editing an existing document" — "Overview" for the lead-in before any "##"
	Text    string // the section's full text, heading line included — what gets embedded, and what Search returns
}

// chunkSkill splits one skill's raw markdown body into Chunks, cutting
// at every "## " line. This mirrors the actual structure every file in
// skills/*.md already uses, so no separate chunk-size tuning
// (domain.ChunkText's char budget, needed for arbitrary source code) is
// necessary here — a skill's own section boundaries are already the
// right retrieval granularity.
func chunkSkill(name, body string) []Chunk {
	var (
		chunks  []Chunk
		heading string
		cur     []string
	)
	flush := func() {
		text := strings.TrimSpace(strings.Join(cur, "\n"))
		if text == "" {
			return
		}
		h := heading
		if h == "" {
			h = "Overview"
		}
		chunks = append(chunks, Chunk{Skill: name, Heading: h, Text: text})
	}
	for _, line := range strings.Split(body, "\n") {
		if strings.HasPrefix(line, "## ") {
			flush()
			heading = strings.TrimPrefix(line, "## ")
			cur = []string{line}
			continue
		}
		cur = append(cur, line)
	}
	flush()
	return chunks
}

// allChunks returns every skill's chunks, in skillkit.Names() order.
func allChunks() []Chunk {
	var out []Chunk
	for _, name := range Names() {
		out = append(out, chunkSkill(name, registry[name].Body)...)
	}
	return out
}

// Match is one ranked search result.
type Match struct {
	Chunk
	Score float64 // cosine similarity to the query, higher is more relevant
}

type vectorChunk struct {
	Chunk
	vec []float32
}

// Index is a lazily-built, in-memory semantic index over every skill's
// chunks. The zero value is not usable — construct with NewIndex. An
// Index is safe for concurrent use and is meant to be built once and
// reused for the life of the process (see package doc for why an
// in-memory index, not a persisted one, is the right call here).
type Index struct {
	embedder Embedder

	mu      sync.Mutex
	vectors []vectorChunk // nil until a build has succeeded at least once
}

// NewIndex returns an Index that will embed every skill chunk, once,
// the first time Search is called (not eagerly here) — so a
// deployment that never calls search_office_skills never pays the
// embedding cost at all, and the first real call pays it exactly once.
func NewIndex(embedder Embedder) *Index {
	return &Index{embedder: embedder}
}

// ensureBuilt embeds every chunk if that hasn't already succeeded.
// Deliberately NOT a sync.Once: if vLLM is unreachable on the very
// first call, that failure must not be cached forever — the next
// Search should try again rather than fail for the rest of the
// process's life over one transient outage. Concurrent callers during
// a build simply serialize on the mutex rather than building twice.
func (ix *Index) ensureBuilt(ctx context.Context) error {
	ix.mu.Lock()
	defer ix.mu.Unlock()
	if ix.vectors != nil {
		return nil
	}
	chunks := allChunks()
	vectors := make([]vectorChunk, 0, len(chunks))
	for _, c := range chunks {
		vec, err := ix.embedder.Embed(ctx, c.Text)
		if err != nil {
			return fmt.Errorf("skillkit: embed %s/%s: %w", c.Skill, c.Heading, err)
		}
		vectors = append(vectors, vectorChunk{Chunk: c, vec: vec})
	}
	ix.vectors = vectors
	return nil
}

// Search embeds query and returns the topK chunks whose own embedding
// is most cosine-similar to it, best first. skillFilter, if non-empty,
// restricts the search to one skill's chunks (matched the same way
// Read does: case-insensitive against the skill name). topK<=0 defaults
// to 4.
func (ix *Index) Search(ctx context.Context, query string, topK int, skillFilter string) ([]Match, error) {
	if topK <= 0 {
		topK = 4
	}
	if err := ix.ensureBuilt(ctx); err != nil {
		return nil, err
	}
	qvec, err := ix.embedder.Embed(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("skillkit: embed query: %w", err)
	}

	ix.mu.Lock()
	vectors := ix.vectors
	ix.mu.Unlock()

	matches := make([]Match, 0, len(vectors))
	for _, v := range vectors {
		if skillFilter != "" && !strings.EqualFold(v.Skill, skillFilter) {
			continue
		}
		matches = append(matches, Match{Chunk: v.Chunk, Score: cosineSimilarity(qvec, v.vec)})
	}
	sort.Slice(matches, func(i, j int) bool { return matches[i].Score > matches[j].Score })
	if len(matches) > topK {
		matches = matches[:topK]
	}
	return matches, nil
}

// cosineSimilarity is the standard cosine of the angle between two
// vectors, in [-1, 1] for any nonzero inputs (in practice embeddings
// from one model land in a narrower positive range). Defensively walks
// only the shorter of the two lengths and returns 0 for a zero vector,
// rather than panicking or dividing by zero, though in normal operation
// every vector here comes from the same embedding model and is the
// same length.
func cosineSimilarity(a, b []float32) float64 {
	n := len(a)
	if len(b) < n {
		n = len(b)
	}
	var dot, normA, normB float64
	for i := 0; i < n; i++ {
		dot += float64(a[i]) * float64(b[i])
		normA += float64(a[i]) * float64(a[i])
		normB += float64(b[i]) * float64(b[i])
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
}