package embedding

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

// Hit is a single keyword search match: the indexed document id and its
// relevance score. Higher scores are more relevant.
type Hit[ID comparable] struct {
	ID    ID
	Score float64
}

// Index is a simple in-memory inverted index for keyword full-text search (FTS).
//
// It tokenizes each added document into terms and keeps a postings list
// (term → doc id → term frequency). Search scores each matching document with
// TF-IDF: the sum over query terms of tf × idf, where idf = ln(1 + N/df). This
// is the keyword-level retrieval counterpart to the semantic Embedder vectors;
// both coexist (see DevDocs/design/04-运行时.md §7).
//
// The zero value is not ready to use; build one with NewIndex. It is safe for
// concurrent use.
type Index[ID comparable] struct {
	mu       sync.RWMutex
	postings map[string]map[ID]int // term -> doc id -> term frequency
	docLen   map[ID]int            // doc id -> term count
	docTerms map[ID][]string       // doc id -> indexed terms, for re-indexing
	docCount int                   // number of distinct documents
}

// NewIndex creates an empty inverted index.
func NewIndex[ID comparable]() *Index[ID] {
	return &Index[ID]{
		postings: make(map[string]map[ID]int),
		docLen:   make(map[ID]int),
		docTerms: make(map[ID][]string),
	}
}

// Add indexes text under id. Re-adding an id replaces its previous term
// contribution, so the index stays consistent with the latest text for that id.
func (ix *Index[ID]) Add(text string, id ID) {
	terms := tokenize(text)
	freq := make(map[string]int, len(terms))
	for _, t := range terms {
		freq[t]++
	}

	ix.mu.Lock()
	defer ix.mu.Unlock()

	if _, exists := ix.docLen[id]; exists {
		ix.removeDocLocked(id)
	}
	for t, f := range freq {
		if ix.postings[t] == nil {
			ix.postings[t] = make(map[ID]int)
		}
		ix.postings[t][id] = f
	}
	ix.docTerms[id] = terms
	ix.docLen[id] = len(terms)
	ix.docCount++
}

// Search returns every document matching at least one query term, ranked by
// descending relevance score. Ties are broken by id for determinism. An empty
// or whitespace-only query returns no hits.
func (ix *Index[ID]) Search(query string) []Hit[ID] {
	terms := tokenize(query)
	if len(terms) == 0 {
		return []Hit[ID]{}
	}

	ix.mu.RLock()
	defer ix.mu.RUnlock()

	scores := make(map[ID]float64)
	n := ix.docCount
	for _, t := range terms {
		df, ok := ix.postings[t]
		if !ok {
			continue
		}
		idf := math.Log(1 + float64(n)/float64(len(df)))
		for id, tf := range df {
			scores[id] += float64(tf) * idf
		}
	}

	hits := make([]Hit[ID], 0, len(scores))
	for id, s := range scores {
		hits = append(hits, Hit[ID]{ID: id, Score: s})
	}
	sort.Slice(hits, func(i, j int) bool {
		if hits[i].Score != hits[j].Score {
			return hits[i].Score > hits[j].Score
		}
		return fmt.Sprint(hits[i].ID) < fmt.Sprint(hits[j].ID)
	})
	return hits
}

// removeDocLocked removes the previous postings contribution of an existing doc
// before it is re-indexed. Callers must hold ix.mu.
func (ix *Index[ID]) removeDocLocked(id ID) {
	for _, t := range ix.docTerms[id] {
		df := ix.postings[t]
		if df != nil {
			delete(df, id)
			if len(df) == 0 {
				delete(ix.postings, t)
			}
		}
	}
	delete(ix.docTerms, id)
	delete(ix.docLen, id)
	ix.docCount--
}

// tokenize splits text into lowercased word terms. Non-letter and non-digit
// runes act as separators; empty and repeated terms are handled by the caller.
func tokenize(s string) []string {
	return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool {
		return !unicode.IsLetter(r) && !unicode.IsDigit(r)
	})
}