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)
	}
}