package datastore

import (
	"context"
	"database/sql"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/semantic"
	"github.com/mem9-ai/drive9/pkg/tenant/schema"
)

func TestUpdateFileEmbedding(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		Revision:    2,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}

	vec := testEmbeddingVector(0.1, 0.2, 0.3)
	updated, err := s.UpdateFileEmbedding(context.Background(), "f1", 2, vec)
	if err != nil {
		t.Fatal(err)
	}
	if !updated {
		t.Fatal("expected update to succeed")
	}
	requireEmbeddingRevision(t, mustFile(t, s, "f1").EmbeddingRevision, 2)

	var raw sql.NullString
	if err := s.DB().QueryRow(`SELECT embedding FROM files WHERE file_id = ?`, "f1").Scan(&raw); err != nil {
		t.Fatal(err)
	}
	got := parseStoredVector(t, raw.String)
	wantPrefix := []float32{0.1, 0.2, 0.3}
	if len(got) < len(wantPrefix) {
		t.Fatalf("embedding=%q parsed=%v, want at least %d components", raw.String, got, len(wantPrefix))
	}
	for i, want := range wantPrefix {
		if d := got[i] - want; d > 1e-4 || d < -1e-4 {
			t.Fatalf("embedding[%d]=%v from %q, want %v", i, got[i], raw.String, want)
		}
	}
}

func TestUpdateFileEmbeddingForGenerationStoresRevisionAndGeneration(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID: "f-generation", StorageType: StorageDB9, StorageRef: "/blobs/f-generation",
		Revision: 2, Status: StatusConfirmed, CreatedAt: now, ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	updated, err := s.UpdateFileEmbeddingForGeneration(context.Background(), "f-generation", 2, 7, testEmbeddingVector(0.4))
	if err != nil || !updated {
		t.Fatalf("UpdateFileEmbeddingForGeneration updated=%v err=%v", updated, err)
	}
	var revision int64
	var generation uint64
	if err := s.DB().QueryRow(`SELECT embedding_revision, embedding_generation FROM semantic WHERE inode_id = ?`, "f-generation").Scan(&revision, &generation); err != nil {
		t.Fatal(err)
	}
	if revision != 2 || generation != 7 {
		t.Fatalf("revision=%d generation=%d, want 2/7", revision, generation)
	}
}

func TestEmbeddingTaskWritebackRejectsSupersededLease(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	if err := s.InsertFile(ctx, &File{
		FileID: "f-superseded", StorageType: StorageDB9, StorageRef: "/blobs/f-superseded",
		Revision: 1, Status: StatusConfirmed, ContentText: "content", Description: "description",
		CreatedAt: now, ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	oldTask := newSemanticTask("task-old-generation", "f-superseded", 1, now, now)
	oldTask.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, oldTask); err != nil {
		t.Fatal(err)
	}
	oldClaim, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim old generation found=%v err=%v", found, err)
	}
	newTask := newSemanticTask("task-new-generation", "f-superseded", 1, now.Add(2*time.Second), now.Add(2*time.Second))
	newTask.PayloadJSON = []byte(`{"generation":2}`)
	if queued, err := s.EnsureSemanticTaskQueued(ctx, newTask); err != nil || !queued {
		t.Fatalf("supersede old generation queued=%v err=%v", queued, err)
	}
	newClaim, found, err := s.ClaimSemanticTask(ctx, now.Add(3*time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim new generation found=%v err=%v", found, err)
	}

	if updated, err := s.UpdateFileEmbeddingForTask(ctx, "f-superseded", 1, 2, newClaim.TaskID, newClaim.Receipt, "content", testEmbeddingVector(0.2)); err != nil || !updated {
		t.Fatalf("write current content embedding updated=%v err=%v", updated, err)
	}
	if updated, err := s.UpdateFileDescriptionEmbeddingForTask(ctx, "f-superseded", 1, 2, newClaim.TaskID, newClaim.Receipt, "description", testEmbeddingVector(0.3)); err != nil || !updated {
		t.Fatalf("write current description embedding updated=%v err=%v", updated, err)
	}
	if updated, err := s.UpdateFileEmbeddingForTask(ctx, "f-superseded", 1, 1, oldClaim.TaskID, oldClaim.Receipt, "content", testEmbeddingVector(0.1)); err != nil || updated {
		t.Fatalf("stale content writeback updated=%v err=%v, want ownership miss", updated, err)
	}
	if updated, err := s.UpdateFileDescriptionEmbeddingForTask(ctx, "f-superseded", 1, 1, oldClaim.TaskID, oldClaim.Receipt, "description", testEmbeddingVector(0.1)); err != nil || updated {
		t.Fatalf("stale description writeback updated=%v err=%v, want ownership miss", updated, err)
	}

	var contentGeneration, descriptionGeneration uint64
	if err := s.DB().QueryRow(`SELECT embedding_generation, description_embedding_generation
		FROM semantic WHERE inode_id = ?`, "f-superseded").Scan(&contentGeneration, &descriptionGeneration); err != nil {
		t.Fatal(err)
	}
	if contentGeneration != 2 || descriptionGeneration != 2 {
		t.Fatalf("generations=%d/%d, want 2/2", contentGeneration, descriptionGeneration)
	}
}

func TestEmbeddingTaskWritebackRejectsChangedContent(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	if err := s.InsertFile(ctx, &File{
		FileID: "f-content-changed", StorageType: StorageDB9, StorageRef: "/blobs/f-content-changed",
		Revision: 1, Status: StatusConfirmed, ContentText: "old transcript",
		CreatedAt: now, ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	task := newSemanticTask("task-content-changed", "f-content-changed", 1, now, now)
	task.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
		t.Fatal(err)
	}
	claim, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	if updated, err := s.UpdateFileSearchText(ctx, "f-content-changed", 1, ""); err != nil || !updated {
		t.Fatalf("clear content updated=%v err=%v", updated, err)
	}

	updated, err := s.UpdateFileEmbeddingForTask(
		ctx, "f-content-changed", 1, 1, claim.TaskID, claim.Receipt, "old transcript", testEmbeddingVector(0.4),
	)
	if err != nil {
		t.Fatal(err)
	}
	if updated {
		t.Fatal("stale content embedding writeback succeeded")
	}
	file := mustFile(t, s, "f-content-changed")
	if file.ContentText != "" || file.EmbeddingRevision != nil {
		t.Fatalf("file content=%q embedding_revision=%v, want empty content and no embedding", file.ContentText, file.EmbeddingRevision)
	}
}

// Description can change without a revision bump (setmeta), so the task
// writeback must skip when the stored description no longer matches the text
// that was sent to the embedding provider.
func TestEmbeddingTaskWritebackRejectsChangedDescription(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	if err := s.InsertFile(ctx, &File{
		FileID: "f-desc-changed", StorageType: StorageDB9, StorageRef: "/blobs/f-desc-changed",
		Revision: 1, Status: StatusConfirmed, Description: "old description",
		CreatedAt: now, ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	task := newSemanticTask("task-desc-changed", "f-desc-changed", 1, now, now)
	task.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
		t.Fatal(err)
	}
	claim, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	if err := s.InTx(ctx, func(tx *sql.Tx) error {
		return s.UpdateFileDescriptionTx(tx, "f-desc-changed", "new description")
	}); err != nil {
		t.Fatalf("setmeta description: %v", err)
	}

	updated, err := s.UpdateFileDescriptionEmbeddingForTask(
		ctx, "f-desc-changed", 1, 1, claim.TaskID, claim.Receipt, "old description", testEmbeddingVector(0.4),
	)
	if err != nil {
		t.Fatal(err)
	}
	if updated {
		t.Fatal("stale description embedding writeback succeeded")
	}
	file := mustFile(t, s, "f-desc-changed")
	if file.Description != "new description" || file.DescriptionEmbeddingRevision != nil {
		t.Fatalf("file description=%q description_embedding_revision=%v, want new description and no embedding",
			file.Description, file.DescriptionEmbeddingRevision)
	}
}

func parseStoredVector(t *testing.T, raw string) []float32 {
	t.Helper()
	s := strings.TrimSpace(raw)
	s = strings.TrimPrefix(s, "[")
	s = strings.TrimSuffix(s, "]")
	if s == "" {
		return nil
	}
	parts := strings.Split(s, ",")
	out := make([]float32, 0, len(parts))
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		f, err := strconv.ParseFloat(p, 32)
		if err != nil {
			t.Fatalf("parse vector component %q from %q: %v", p, raw, err)
		}
		out = append(out, float32(f))
	}
	return out
}

func TestUpdateFileEmbeddingSkipsStaleRevision(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		Revision:    3,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}

	updated, err := s.UpdateFileEmbedding(context.Background(), "f1", 2, testEmbeddingVector(0.1, 0.2))
	if err != nil {
		t.Fatal(err)
	}
	if updated {
		t.Fatal("stale revision should not update embedding")
	}
	if got := mustFile(t, s, "f1").EmbeddingRevision; got != nil {
		t.Fatalf("embedding revision=%v, want nil", *got)
	}
}

func mustFile(t *testing.T, s *Store, fileID string) *File {
	t.Helper()
	f, err := s.GetFile(context.Background(), fileID)
	if err != nil {
		t.Fatalf("get file %s: %v", fileID, err)
	}
	return f
}

func testEmbeddingVector(prefix ...float32) []float32 {
	vec := make([]float32, schema.TiDBAutoEmbeddingDimensions)
	copy(vec, prefix)
	return vec
}