package memory

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

// NonRecallPrefix — rows whose id starts with this are excluded from Search by
// every backend. Speaker voiceprints live in the same vector table as notes and
// facts (internal/speaker writes them under this prefix), and they are not
// recall material: a voiceprint has no text to read back and surfacing one as a
// note hit leaks a name attached to a biometric.
//
// Reading them through Catalog.ByPrefix was documented as what keeps them out
// of recall. It is not. It controls how speaker code reads its own rows and
// says nothing about what Search scores. What actually kept them out was that
// cosine returns 0 on a width mismatch, so a 192-dim voiceprint scored 0
// against a 384-dim query. That is a coincidence of two model choices — some
// x-vector exports are 384-dim — and not an invariant. This is the invariant.
const NonRecallPrefix = "speaker:"

// Result is a single search hit.
type Result struct {
	ID    string
	Score float64
	Meta  map[string]string
}

// Store is a vector memory interface.
type Store interface {
	Insert(ctx context.Context, id string, vec []float32, meta map[string]string) error
	Search(ctx context.Context, vec []float32, topK int) ([]Result, error)
}

// Record is a stored vector read back whole — id, vector and metadata — as
// opposed to Result, which is a search hit and carries a score instead of the
// vector.
type Record struct {
	ID   string
	Vec  []float32
	Meta map[string]string
}

// Catalog is a Store that can also be enumerated by id prefix and deleted from.
//
// Search is not enough for every user of the vector table. Speaker profiles
// (internal/speaker) need to list exactly their own rows without scoring
// anything, because listing enrolled voices is not a similarity question, and
// they need Delete because a voiceprint is data about a person and "forget this
// voice" has to actually remove it. Note and fact recall use plain Store and are
// unaffected.
type Catalog interface {
	Store
	// ByPrefix returns every row whose id starts with prefix, in no particular
	// order. An empty prefix returns everything.
	ByPrefix(ctx context.Context, prefix string) ([]Record, error)
	// Delete removes one row by id. Deleting a row that is not there is not an
	// error: the caller asked for it to be gone and it is gone.
	Delete(ctx context.Context, id string) error
}

// item is a single stored vector with metadata.
type item struct {
	id   string
	vec  []float32
	meta map[string]string
}

// InMemoryStore implements Store with cosine similarity search.
type InMemoryStore struct {
	mu    sync.RWMutex
	items []item
}

// compile-time check: InMemoryStore satisfies Catalog.
var _ Catalog = (*InMemoryStore)(nil)

func NewInMemoryStore() *InMemoryStore {
	return &InMemoryStore{}
}

// Insert upserts by id, matching the persistent store.MemoryStore: a repeated
// id replaces the prior row rather than accumulating a second copy. Re-indexing
// a note is an update, and re-enrolling a voice must replace the old voiceprint
// rather than leave it searchable.
func (s *InMemoryStore) Insert(_ context.Context, id string, vec []float32, meta map[string]string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.items {
		if s.items[i].id == id {
			s.items[i] = item{id: id, vec: vec, meta: meta}
			return nil
		}
	}
	s.items = append(s.items, item{id: id, vec: vec, meta: meta})
	return nil
}

// ByPrefix implements Catalog.
func (s *InMemoryStore) ByPrefix(_ context.Context, prefix string) ([]Record, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	var out []Record
	for _, it := range s.items {
		if !strings.HasPrefix(it.id, prefix) {
			continue
		}
		// Copy the metadata too. Returning it.meta by reference let a caller
		// mutating the returned map edit the stored row, and the persistent
		// backend unmarshals fresh, so the two disagreed.
		meta := make(map[string]string, len(it.meta))
		for k, v := range it.meta {
			meta[k] = v
		}
		out = append(out, Record{ID: it.id, Vec: append([]float32(nil), it.vec...), Meta: meta})
	}
	return out, nil
}

// Delete implements Catalog.
func (s *InMemoryStore) Delete(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.items {
		if s.items[i].id == id {
			s.items = append(s.items[:i], s.items[i+1:]...)
			return nil
		}
	}
	return nil
}

func (s *InMemoryStore) Search(_ context.Context, vec []float32, topK int) ([]Result, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if topK <= 0 {
		topK = 10
	}

	type scored struct {
		id    string
		score float64
		meta  map[string]string
	}

	scores := make([]scored, 0, len(s.items))
	for _, it := range s.items {
		if strings.HasPrefix(it.id, NonRecallPrefix) {
			continue
		}
		score := cosine(vec, it.vec)
		scores = append(scores, scored{id: it.id, score: score, meta: it.meta})
	}

	sort.Slice(scores, func(i, j int) bool {
		return scores[i].score > scores[j].score // descending
	})

	if topK > len(scores) {
		topK = len(scores)
	}

	out := make([]Result, topK)
	for i := 0; i < topK; i++ {
		out[i] = Result{
			ID:    scores[i].id,
			Score: scores[i].score,
			Meta:  scores[i].meta,
		}
	}
	return out, nil
}

// cosine similarity (dot product, assumes L2-normalized vectors).
func cosine(a, b []float32) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var dot float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
	}
	return dot
}