package embedding

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

func TestNGramEmbedder_Dim(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	if got := e.Dim(); got != 128 {
		t.Fatalf("Dim() = %d, want 128", got)
	}
}

func TestNGramEmbedder_DefaultDim(t *testing.T) {
	t.Parallel()

	e := NewNGram(0)
	if got := e.Dim(); got != defaultNGramDim {
		t.Fatalf("Dim() = %d, want defaultNGramDim (%d)", got, defaultNGramDim)
	}
}

func TestNGramEmbedder_VectorLengthMatchesDim(t *testing.T) {
	t.Parallel()

	e := NewNGram(64)
	vectors, err := e.Embed(context.Background(), []string{"alpha", "beta gamma"})
	if err != nil {
		t.Fatalf("Embed() error: %v", err)
	}
	if len(vectors) != 2 {
		t.Fatalf("got %d vectors, want 2", len(vectors))
	}
	for i, v := range vectors {
		if got := len(v); got != e.Dim() {
			t.Errorf("vector %d length = %d, want %d", i, got, e.Dim())
		}
	}
}

func TestNGramEmbedder_Deterministic(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	ctx := context.Background()
	texts := []string{"hello world", "the quick brown fox jumps", "Golang embedding"}

	first, err := e.Embed(ctx, texts)
	if err != nil {
		t.Fatalf("Embed() error: %v", err)
	}
	second, err := e.Embed(ctx, texts)
	if err != nil {
		t.Fatalf("Embed() error: %v", err)
	}
	if len(first) != len(second) {
		t.Fatalf("got %d vectors on first call, %d on second", len(first), len(second))
	}
	for i := range first {
		if !equalVector(first[i], second[i]) {
			t.Fatalf("vector %d differs between calls", i)
		}
	}
}

func TestNGramEmbedder_Normalized(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	vectors, err := e.Embed(context.Background(), []string{
		"a",
		"the quick brown fox jumps over the lazy dog",
		"golang fts embedding test",
	})
	if err != nil {
		t.Fatalf("Embed() error: %v", err)
	}
	for i, v := range vectors {
		norm := l2Norm(v)
		if math.Abs(norm-1) > 1e-6 {
			t.Errorf("vector %d L2 norm = %f, want 1", i, norm)
		}
	}
}

func TestNGramEmbedder_DifferentTextsDiffer(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	vectors, err := e.Embed(context.Background(), []string{"dog", "cat", "astronaut"})
	if err != nil {
		t.Fatalf("Embed() error: %v", err)
	}
	if len(vectors) != 3 {
		t.Fatalf("got %d vectors, want 3", len(vectors))
	}
	for i := 0; i < len(vectors); i++ {
		if l2Norm(vectors[i]) == 0 {
			t.Errorf("vector %d is a zero vector", i)
		}
		for j := i + 1; j < len(vectors); j++ {
			if equalVector(vectors[i], vectors[j]) {
				t.Errorf("vectors %d and %d are identical", i, j)
			}
		}
	}
}

func TestNGramEmbedder_EmptyInput(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	vectors, err := e.Embed(context.Background(), nil)
	if err != nil {
		t.Fatalf("Embed(nil) error: %v", err)
	}
	if len(vectors) != 0 {
		t.Errorf("got %d vectors, want 0", len(vectors))
	}
}

func TestNGramEmbedder_CanceledContext(t *testing.T) {
	t.Parallel()

	e := NewNGram(128)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := e.Embed(ctx, []string{"hello"}); err == nil {
		t.Fatal("expected error on canceled context, got nil")
	}
}

func equalVector(a, b Vector) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func l2Norm(v Vector) float64 {
	var sum float64
	for _, x := range v {
		sum += float64(x) * float64(x)
	}
	return math.Sqrt(sum)
}