//go:build failpoint

package server

import (
	"context"
	"crypto/rand"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/c4pt0r/agfs/agfs-server/pkg/filesystem"
	mysql "github.com/go-sql-driver/mysql"
	"github.com/pingcap/failpoint"

	"github.com/mem9-ai/drive9/pkg/backend"
	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/encrypt"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/s3client"
	"github.com/mem9-ai/drive9/pkg/semantic"
	"github.com/mem9-ai/drive9/pkg/tenant"
	"github.com/mem9-ai/drive9/pkg/testtidb"
)

const semanticWorkerBeforeFinalizeFailpoint = "github.com/mem9-ai/drive9/pkg/server/semanticWorkerBeforeFinalize"

// gatedImgExtractor blocks the first extraction call until released, so the
// test can delete the tenant while a claimed img_extract_text task is inside
// the provider call.
type gatedImgExtractor struct {
	started chan struct{}
	release chan struct{}

	mu    sync.Mutex
	calls int
}

func (e *gatedImgExtractor) ExtractImageText(ctx context.Context, _ backend.ImageExtractRequest) (string, backend.ImageExtractUsage, error) {
	e.mu.Lock()
	e.calls++
	call := e.calls
	e.mu.Unlock()
	if call == 1 {
		select {
		case e.started <- struct{}{}:
		default:
		}
		select {
		case <-e.release:
		case <-ctx.Done():
			return "", backend.ImageExtractUsage{}, ctx.Err()
		}
		return "stale caption", backend.ImageExtractUsage{}, nil
	}
	return "caption", backend.ImageExtractUsage{}, nil
}

func (e *gatedImgExtractor) callCount() int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.calls
}

// TestImgExtractTaskTenantDeleteRace covers the resolution/delete race: a
// claimed img_extract_text task resolves the tenant runtime, blocks inside
// the extractor, and the tenant is deleted mid-flight (the meta delete path
// wipes the extract config rows). After the extractor releases, the task must
// finalize with a failed/no-op writeback, and no new claim for the deleted
// tenant may succeed.
func TestImgExtractTaskTenantDeleteRace(t *testing.T) {
	ctx := context.Background()
	initServerTenantSchema(t, testDSN)
	store, err := datastore.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	testtidb.ResetDB(t, store.DB())
	t.Cleanup(func() { _ = store.Close() })

	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())

	parsed, err := mysql.ParseDSN(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	host, port := "127.0.0.1", 3306
	if parsed.Addr != "" {
		h, p, _ := strings.Cut(parsed.Addr, ":")
		if h != "" {
			host = h
		}
		if p != "" {
			_, _ = fmt.Sscanf(p, "%d", &port)
		}
	}
	master := make([]byte, 32)
	if _, err := rand.Read(master); err != nil {
		t.Fatal(err)
	}
	enc, err := encrypt.NewLocalAESEncryptor(master)
	if err != nil {
		t.Fatal(err)
	}
	pool := tenant.NewPool(tenant.PoolConfig{S3Dir: mustTempDir(t), PublicURL: "http://localhost", SkipTiDBSchemaCheck: true}, enc)
	pool.SetMetaStore(metaStore)
	t.Cleanup(pool.Close)

	tenantID := "t-img-delete-race"
	passCipher, err := pool.Encrypt(ctx, []byte(parsed.Passwd))
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC()
	if err := metaStore.InsertTenant(ctx, &meta.Tenant{
		ID:               tenantID,
		Status:           meta.TenantActive,
		Kind:             meta.TenantKindLive,
		DBHost:           host,
		DBPort:           port,
		DBUser:           parsed.User,
		DBPasswordCipher: passCipher,
		DBName:           parsed.DBName,
		Provider:         tenant.ProviderTiDBCloudNative,
		ClusterID:        "cluster-a",
		BranchID:         "branch-a",
		SchemaVersion:    1,
		CreatedAt:        now,
		UpdatedAt:        now,
	}); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.UpsertTenantExtractConfig(ctx, &meta.TenantExtractConfig{
		TenantID:     tenantID,
		MediaType:    meta.ExtractMediaTypeImage,
		Enabled:      true,
		APIBase:      "https://extractor.example.com",
		APIKeyCipher: []byte("cipher"),
		Model:        "vision-model-1",
	}); err != nil {
		t.Fatal(err)
	}

	extractor := &gatedImgExtractor{started: make(chan struct{}, 1), release: make(chan struct{})}
	// Nanosecond TTL so post-delete re-resolution never serves the pre-delete
	// cache entry.
	resolver := newExtractConfigResolver(metaStore, &fakeExtractDecryptor{plaintext: "k"}, nil, time.Nanosecond)
	t.Cleanup(resolver.Close)
	resolver.newOpenAIExtractor = func(backend.OpenAIImageTextExtractorConfig) (backend.ImageTextExtractor, error) {
		return extractor, nil
	}

	s3c, err := s3client.NewLocal(mustTempDir(t), "http://localhost:9091/s3")
	if err != nil {
		t.Fatal(err)
	}
	b, err := backend.NewWithS3ModeAndOptions(store, s3c, true, backend.Options{
		AsyncImageExtract: backend.AsyncImageExtractOptions{Capability: true},
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := b.Write("/img/race.png", fakeImagePNG, 0, filesystem.WriteFlagCreate); err != nil {
		t.Fatal(err)
	}

	m := newTenantWorkerManager(nil, metaStore, pool, nil, TenantWorkerOptions{}, 0)
	if m == nil {
		t.Fatal("expected viable worker manager")
	}
	m.setExtractRuntimeDeps(resolver, nil)

	// Capture the finalize decision for the img task at the state-transition
	// boundary (before ack/retry/dead-letter is applied).
	type finalizeEvent struct{ action, result string }
	events := make(chan finalizeEvent, 4)
	if err := failpoint.EnableCall(semanticWorkerBeforeFinalizeFailpoint,
		func(_ string, _ *datastore.Store, task *semantic.Task, action, _, result string) {
			if task != nil && task.TaskType == semantic.TaskTypeImgExtractText {
				events <- finalizeEvent{action: action, result: result}
			}
		}); err != nil {
		t.Fatalf("enable failpoint: %v", err)
	}
	t.Cleanup(func() { _ = failpoint.Disable(semanticWorkerBeforeFinalizeFailpoint) })

	target := &tenantTarget{
		tenantID:         tenantID,
		backend:          b,
		store:            b.Store(),
		allowedTaskTypes: []semantic.TaskType{semantic.TaskTypeImgExtractText},
		release:          func() {},
	}
	claimed := make(chan bool, 1)
	go func() { ok, _ := m.claimAndProcessOne(ctx, target); claimed <- ok }()

	select {
	case <-extractor.started:
	case <-time.After(10 * time.Second):
		t.Fatal("extractor did not start")
	}

	// Delete the tenant mid-extraction: the meta delete path wipes the extract
	// config rows; the tenant file rows vanish with the tenant data.
	if err := metaStore.MarkTenantDeleted(ctx, tenantID); err != nil {
		t.Fatal(err)
	}
	for _, q := range []string{
		"DELETE FROM file_nodes",
		"DELETE FROM file_tags",
		"DELETE FROM files",
		"DELETE FROM inodes",
		"DELETE FROM contents",
	} {
		if _, err := store.DB().ExecContext(ctx, q); err != nil {
			t.Fatalf("simulate tenant data teardown: %v", err)
		}
	}
	close(extractor.release)

	select {
	case ok := <-claimed:
		if !ok {
			t.Fatal("claimAndProcessOne returned false, want claimed task")
		}
	case <-time.After(10 * time.Second):
		t.Fatal("claimAndProcessOne did not return")
	}

	// The task finalized with a no-op writeback (stale: the revision-gated
	// update matched no row after the tenant data was wiped), acked exactly
	// once, and the provider was called exactly once.
	select {
	case ev := <-events:
		if ev.action != string(semanticTaskActionAck) || ev.result != string(backend.ImageExtractResultStale) {
			t.Fatalf("finalize = %+v, want ack/%s", ev, backend.ImageExtractResultStale)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("finalize failpoint did not fire")
	}
	if extractor.callCount() != 1 {
		t.Fatalf("extractor calls = %d, want 1", extractor.callCount())
	}
	assertSemanticTaskCount(t, store, string(semantic.TaskSucceeded), 1)

	// No new claim succeeds for the deleted tenant: the scheduler rejects the
	// kick, and re-resolution confirms the config rows are gone.
	if _, ok, err := m.kickRef(ctx, tenantID, WorkSemantic); err != nil || ok {
		t.Fatalf("kickRef for deleted tenant: ok=%v err=%v, want rejected", ok, err)
	}
	res, err := resolver.ResolveImage(ctx, tenantID)
	if err != nil {
		t.Fatalf("ResolveImage after delete: %v", err)
	}
	if res.state != extractResolveNotConfigured {
		t.Fatalf("post-delete resolve state = %v, want not_configured (config rows wiped)", res.state)
	}
}

func TestTerminalMediaConfigLeaseLossPreservesQueuedPeers(t *testing.T) {
	ctx := context.Background()
	initServerTenantSchema(t, testDSN)
	store, err := datastore.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	testtidb.ResetDB(t, store.DB())
	t.Cleanup(func() { _ = store.Close() })

	now := time.Now().UTC()
	for i := range 2 {
		task := &semantic.Task{
			TaskID: fmt.Sprintf("lease-loss-terminal-%d", i), TaskType: semantic.TaskTypeImgExtractText,
			ResourceID: fmt.Sprintf("lease-loss-file-%d", i), ResourceVersion: 1,
			Status: semantic.TaskQueued, MaxAttempts: 5, AvailableAt: now,
			CreatedAt: now.Add(time.Duration(i) * time.Millisecond), UpdatedAt: now,
		}
		if _, err := store.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}
	claimed, found, err := store.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeImgExtractText)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	resolver := newExtractConfigResolver(
		&fakeTenantExtractConfigReader{rows: map[string]meta.TenantExtractConfig{}},
		&fakeExtractDecryptor{}, nil, time.Minute,
	)
	t.Cleanup(resolver.Close)
	m := &tenantWorkerManager{extractResolver: resolver}
	target := &tenantTarget{tenantID: "lease-loss-terminal", store: store}

	var expireErr error
	if err := failpoint.EnableCall(semanticWorkerBeforeFinalizeFailpoint,
		func(_ string, _ *datastore.Store, task *semantic.Task, _, _, _ string) {
			if task == nil || task.TaskID != claimed.TaskID {
				return
			}
			_, expireErr = store.DB().ExecContext(ctx,
				`UPDATE semantic_tasks SET lease_until = ? WHERE task_id = ?`,
				now.Add(-time.Minute), task.TaskID,
			)
		}); err != nil {
		t.Fatalf("enable failpoint: %v", err)
	}
	t.Cleanup(func() { _ = failpoint.Disable(semanticWorkerBeforeFinalizeFailpoint) })

	m.processTask(ctx, target, claimed)
	if expireErr != nil {
		t.Fatalf("expire claimed lease: %v", expireErr)
	}
	var queued int
	if err := store.DB().QueryRowContext(ctx,
		`SELECT COUNT(*) FROM semantic_tasks WHERE task_type = ? AND status = ?`,
		semantic.TaskTypeImgExtractText, semantic.TaskQueued,
	).Scan(&queued); err != nil {
		t.Fatal(err)
	}
	if queued != 1 {
		t.Fatalf("queued peers=%d, want 1 when claimed task loses ownership before finalize", queued)
	}
}