package memory

import (
	"context"
	"fmt"
	"math"
	"testing"
)

func TestInsertAndSearch(t *testing.T) {
	s := NewInMemoryStore()
	ctx := context.Background()

	// Insert a few vectors
	s.Insert(ctx, "doc1", []float32{1, 0, 0}, nil)
	s.Insert(ctx, "doc2", []float32{0, 1, 0}, nil)
	s.Insert(ctx, "doc3", []float32{0, 0, 1}, map[string]string{"source": "note"})

	// Search for something close to doc1
	results, err := s.Search(ctx, []float32{0.9, 0.1, 0}, 5)
	if err != nil {
		t.Fatalf("Search: %v", err)
	}
	if len(results) != 3 {
		t.Fatalf("expected 3 results, got %d", len(results))
	}
	if results[0].ID != "doc1" {
		t.Errorf("nearest should be doc1, got %s", results[0].ID)
	}
	if math.Abs(results[0].Score-0.9) > 0.01 {
		t.Errorf("doc1 score should be near 0.9, got %f", results[0].Score)
	}

	// Check metadata preserved
	if results[2].Meta["source"] != "note" {
		t.Errorf("doc3 meta.source = %q, want note", results[2].Meta["source"])
	}
}

func TestTopKTruncation(t *testing.T) {
	s := NewInMemoryStore()
	ctx := context.Background()

	// Distinct ids: Insert upserts by id, so ten rows need ten ids.
	for i := 0; i < 10; i++ {
		s.Insert(ctx, fmt.Sprintf("n%d", i), []float32{float32(i) / 10, 0, 0}, nil)
	}

	results, err := s.Search(ctx, []float32{1, 0, 0}, 3)
	if err != nil {
		t.Fatalf("Search: %v", err)
	}
	if len(results) != 3 {
		t.Fatalf("expected 3 results with topK=3, got %d", len(results))
	}
}

func TestEmptyStore(t *testing.T) {
	s := NewInMemoryStore()
	results, err := s.Search(context.Background(), []float32{1, 0, 0}, 5)
	if err != nil {
		t.Fatalf("Search on empty store: %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("expected 0 results, got %d", len(results))
	}
}

func TestCosineEdgeCases(t *testing.T) {
	if c := cosine(nil, []float32{1}); c != 0 {
		t.Errorf("nil first: expected 0, got %f", c)
	}
	if c := cosine([]float32{1}, nil); c != 0 {
		t.Errorf("nil second: expected 0, got %f", c)
	}
	if c := cosine([]float32{}, []float32{}); c != 0 {
		t.Errorf("empty: expected 0, got %f", c)
	}
	if c := cosine([]float32{1, 2}, []float32{1, 2}); math.Abs(c-5) > 0.001 {
		t.Errorf("dot(1,2;1,2) = %f, want 5", c)
	}
}

// The in-memory backend has to hide voiceprints from recall exactly like the
// persistent one, or a test passing here says nothing about the daemon.
func TestInMemorySearchSkipsVoiceprints(t *testing.T) {
	ctx := context.Background()
	s := NewInMemoryStore()
	if err := s.Insert(ctx, "note:1", []float32{0, 1}, map[string]string{"text": "заметка"}); err != nil {
		t.Fatal(err)
	}
	if err := s.Insert(ctx, NonRecallPrefix+"kami", []float32{1, 0}, map[string]string{"name": "Ками"}); err != nil {
		t.Fatal(err)
	}
	got, err := s.Search(ctx, []float32{1, 0}, 10)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 || got[0].ID != "note:1" {
		t.Fatalf("Search = %+v, want just the note", got)
	}
	recs, err := s.ByPrefix(ctx, NonRecallPrefix)
	if err != nil || len(recs) != 1 {
		t.Fatalf("ByPrefix = %+v, %v; want the voiceprint", recs, err)
	}
}

// ByPrefix hands back a copy of the metadata. It used to return the stored map
// by reference, so a caller editing a returned Record silently edited the row,
// and the persistent backend did not behave that way.
func TestInMemoryByPrefixCopiesMeta(t *testing.T) {
	ctx := context.Background()
	s := NewInMemoryStore()
	if err := s.Insert(ctx, "speaker:kami", []float32{1, 0}, map[string]string{"name": "Ками"}); err != nil {
		t.Fatal(err)
	}
	recs, err := s.ByPrefix(ctx, "speaker:")
	if err != nil || len(recs) != 1 {
		t.Fatalf("ByPrefix = %+v, %v", recs, err)
	}
	recs[0].Meta["name"] = "не Ками"

	again, err := s.ByPrefix(ctx, "speaker:")
	if err != nil {
		t.Fatal(err)
	}
	if again[0].Meta["name"] != "Ками" {
		t.Errorf("the stored row was edited through the returned map: %q", again[0].Meta["name"])
	}
}

// Insert upserts. This is not only a speaker-profile concern: every user of the
// in-memory store used to accumulate a second row for a re-indexed id, and the
// stale copy stayed searchable.
func TestInMemoryInsertUpserts(t *testing.T) {
	ctx := context.Background()
	s := NewInMemoryStore()
	if err := s.Insert(ctx, "note:1", []float32{1, 0}, map[string]string{"text": "старое"}); err != nil {
		t.Fatal(err)
	}
	if err := s.Insert(ctx, "note:1", []float32{0, 1}, map[string]string{"text": "новое"}); err != nil {
		t.Fatal(err)
	}
	got, err := s.Search(ctx, []float32{1, 0}, 10)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 {
		t.Fatalf("Search returned %d rows, want 1 (the old copy is still searchable)", len(got))
	}
	if got[0].Meta["text"] != "новое" {
		t.Errorf("row = %q, want the replacement", got[0].Meta["text"])
	}
}