package datastore

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

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

func TestEmbeddingTaskWritebackSharedScope(t *testing.T) {
	installSharedCoreFSNoLegacy(t)
	ctx := context.Background()
	now := time.Now().UTC()
	storeA := newSharedStore(t, 4400021)
	storeB := newSharedStore(t, 4400022)

	for _, store := range []*Store{storeA, storeB} {
		if err := store.InsertFile(ctx, &File{
			FileID: "shared-file", StorageType: StorageDB9, StorageRef: "/blobs/shared-file",
			Revision: 1, Status: StatusConfirmed, ContentText: "content",
			CreatedAt: now, ConfirmedAt: &now,
		}); err != nil {
			t.Fatal(err)
		}
		task := newSemanticTask("shared-task", "shared-file", 1, now, now)
		task.PayloadJSON = []byte(`{"generation":1}`)
		if _, err := store.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}

	claimAt := now.Add(time.Second)
	claimA, found, err := storeA.ClaimSemanticTask(ctx, claimAt, time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim A found=%v err=%v", found, err)
	}
	claimB, found, err := storeB.ClaimSemanticTask(ctx, claimAt, time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim B found=%v err=%v", found, err)
	}
	if claimA.Receipt == claimB.Receipt {
		t.Fatal("shared task claims unexpectedly reused a receipt")
	}

	updated, err := storeA.UpdateFileEmbeddingForTask(
		ctx, "shared-file", 1, 9, "shared-task", claimB.Receipt, "content", testEmbeddingVector(0.9),
	)
	if err != nil || updated {
		t.Fatalf("cross-scope receipt write updated=%v err=%v, want ownership miss", updated, err)
	}

	for _, tc := range []struct {
		name       string
		store      *Store
		receipt    string
		generation uint64
	}{
		{name: "A", store: storeA, receipt: claimA.Receipt, generation: 11},
		{name: "B", store: storeB, receipt: claimB.Receipt, generation: 22},
	} {
		updated, err := tc.store.UpdateFileEmbeddingForTask(
			ctx, "shared-file", 1, tc.generation, "shared-task", tc.receipt, "content", testEmbeddingVector(float32(tc.generation)),
		)
		if err != nil || !updated {
			t.Fatalf("write %s updated=%v err=%v", tc.name, updated, err)
		}
	}

	for _, tc := range []struct {
		name       string
		fsID       int64
		generation uint64
	}{
		{name: "A", fsID: 4400021, generation: 11},
		{name: "B", fsID: 4400022, generation: 22},
	} {
		var generation sql.NullInt64
		if err := storeA.DB().QueryRowContext(ctx, `SELECT embedding_generation FROM semantic
			WHERE fs_id = ? AND inode_id = ?`, tc.fsID, "shared-file").Scan(&generation); err != nil {
			t.Fatal(err)
		}
		if !generation.Valid || uint64(generation.Int64) != tc.generation {
			t.Fatalf("generation %s=%v, want %d", tc.name, generation, tc.generation)
		}
	}
}