package memory

import (
	"context"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/projectindex"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory/vector"
)

// mockVectorIndex is a scriptable vector.Index for tests. With no fns it emits
// one-hot vectors and returns no search hits.
type mockVectorIndex struct {
	dim      int
	embedFn  func(ctx context.Context, texts []string) ([]vector.Vector, error)
	searchFn func(ctx context.Context, q vector.Vector, topK int) ([]vector.Hit, error)
}

func (m *mockVectorIndex) Dim() int {
	if m.dim > 0 {
		return m.dim
	}
	return 4
}

func (m *mockVectorIndex) Add(_ context.Context, _ string, _ vector.Vector) error { return nil }

func (m *mockVectorIndex) Close() error { return nil }

func (m *mockVectorIndex) Embed(ctx context.Context, texts []string) ([]vector.Vector, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if m.embedFn != nil {
		return m.embedFn(ctx, texts)
	}
	out := make([]vector.Vector, len(texts))
	for i := range texts {
		v := make(vector.Vector, m.Dim())
		v[i%m.Dim()] = 1
		out[i] = v
	}
	return out, nil
}

func (m *mockVectorIndex) Search(ctx context.Context, _ vector.Vector, topK int) ([]vector.Hit, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if m.searchFn != nil {
		return m.searchFn(ctx, nil, topK)
	}
	return []vector.Hit{}, nil
}

func TestRetriever_Retrieve_WithVectorIndex_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	// Both items are in English while the query is Chinese, so the keyword path
	// shares no token and contributes nothing. The ranking is therefore driven
	// entirely by the vector index, proving the persistent vector path is used.
	item1 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she is a morning person",
		Scope: ScopePersona, Kind: KindRelationship,
	})
	item2 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she likes jazz",
		Scope: ScopePersona, Kind: KindRelationship,
	})

	idx := &mockVectorIndex{
		searchFn: func(_ context.Context, _ vector.Vector, _ int) ([]vector.Hit, error) {
			return []vector.Hit{
				{ID: item1.ID, Score: 0.9},
				{ID: item2.ID, Score: 0.2},
			}, nil
		},
	}
	r := NewRetriever(s, WithVectorIndex(idx))

	results, err := r.Retrieve(ctx, "persona.a", "咖啡", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected results with a vector index wired")
	}
	if got := results[0].Item.Content; got != "she is a morning person" {
		t.Fatalf("top result = %q, want the item the vector index ranked highest", got)
	}
	if results[0].Score <= 0 {
		t.Fatalf("top result score = %v, want > 0", results[0].Score)
	}
}

func TestRetriever_Retrieve_WithVectorIndex_PartialHits_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	item1 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she is a morning person",
		Scope: ScopePersona, Kind: KindRelationship,
	})
	item2 := addItem(t, s, MemoryItem{
		PersonaID: "persona.a", Content: "she likes jazz",
		Scope: ScopePersona, Kind: KindRelationship,
	})

	// The vector index only knows about item1; item2 is absent from it, so it
	// carries no semantic score. With a Chinese query that the keyword path
	// cannot match, only item1 is relevant: it must rank at the top with a
	// positive score while item2 scores 0, proving the semantic contribution
	// comes strictly from the vector index.
	idx := &mockVectorIndex{
		searchFn: func(_ context.Context, _ vector.Vector, _ int) ([]vector.Hit, error) {
			return []vector.Hit{{ID: item1.ID, Score: 0.9}}, nil
		},
	}
	r := NewRetriever(s, WithVectorIndex(idx))

	results, err := r.Retrieve(ctx, "persona.a", "咖啡", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("got %d results, want 2", len(results))
	}
	if got := results[0].Item.Content; got != "she is a morning person" {
		t.Fatalf("top result = %q, want the vector-indexed item", got)
	}
	if results[0].Score <= 0 {
		t.Fatalf("top result score = %v, want > 0", results[0].Score)
	}
	// item2 has no vector hit and no keyword overlap, so it scores 0 (matched
	// mode) and sorts below.
	if results[1].Item.ID != item2.ID {
		t.Fatalf("second result id = %q, want item2", results[1].Item.ID)
	}
	if results[1].Score != 0 {
		t.Fatalf("second result score = %v, want 0 (no vector hit)", results[1].Score)
	}
}

func TestRetriever_Retrieve_Source_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	s := newTestStore(t)

	addItem(t, s, MemoryItem{
		PersonaID:          "persona.a",
		Content:            "她喜欢冰美式",
		Tags:               []string{"coffee"},
		SourceConversation: "conv.42",
		Scope:              ScopePersona,
		Kind:               KindRelationship,
	})

	r := NewRetriever(s)
	results, err := r.Retrieve(ctx, "persona.a", "冰美式", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].SourceConversation != "conv.42" {
		t.Fatalf("Result.SourceConversation = %q, want conv.42 (溯源)", results[0].SourceConversation)
	}
	if results[0].Item.SourceConversation != "conv.42" {
		t.Fatalf("Item.SourceConversation = %q, want conv.42", results[0].Item.SourceConversation)
	}
}

func TestRetriever_Retrieve_ProjectLink_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()

	// Separate project-index DB, so it shares no connection with the memory store.
	projDB, err := db.Open(ctx, ":memory:", testLogger())
	if err != nil {
		t.Fatalf("open project db: %v", err)
	}
	t.Cleanup(func() { _ = projDB.Close() })
	ps, err := projectindex.New(ctx, projDB, testLogger())
	if err != nil {
		t.Fatalf("projectindex.New: %v", err)
	}
	if err := ps.PutProject(ctx, projectindex.Project{
		ID: "project.alpha", Title: "Alpha 项目", Desc: "Alpha 项目描述",
	}); err != nil {
		t.Fatalf("PutProject: %v", err)
	}

	ms := newTestStore(t)
	addItem(t, ms, MemoryItem{
		PersonaID: "persona.a",
		Content:   "Alpha 项目的一个事实",
		Scope:     ScopeProject,
		ProjectID: "project.alpha",
		Kind:      KindWorldEvent,
	})

	r := NewRetriever(ms, WithProjectIndex(ps))
	results, err := r.Retrieve(ctx, "persona.a", "Alpha", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].Item.Scope != ScopeProject {
		t.Fatalf("result scope = %q, want project", results[0].Item.Scope)
	}
	if results[0].Project == nil {
		t.Fatal("Result.Project is nil, want the linked project-index entry")
	}
	if results[0].Project.ID != "project.alpha" {
		t.Fatalf("Result.Project.ID = %q, want project.alpha", results[0].Project.ID)
	}
}

func TestRetriever_Retrieve_ProjectLink_NoStore_正常(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	ms := newTestStore(t)

	addItem(t, ms, MemoryItem{
		PersonaID: "persona.a", Content: "Alpha 项目的一个事实",
		Scope: ScopeProject, ProjectID: "project.alpha", Kind: KindWorldEvent,
	})

	// Without a project-index store wired, the hit still returns and Project is
	// simply left unset (no error).
	r := NewRetriever(ms)
	results, err := r.Retrieve(ctx, "persona.a", "Alpha", 10)
	if err != nil {
		t.Fatalf("Retrieve: %v", err)
	}
	if len(results) == 0 {
		t.Fatal("expected a result")
	}
	if results[0].Project != nil {
		t.Fatalf("Result.Project = %+v, want nil without a project-index store", results[0].Project)
	}
}