package tenant

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/go-sql-driver/mysql"

	"github.com/mem9-ai/drive9/pkg/backend"
	"github.com/mem9-ai/drive9/pkg/embedding"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/semantic"
	"github.com/mem9-ai/drive9/pkg/tenant/schema"
	"github.com/mem9-ai/drive9/pkg/testtidb"
)

type sharedDummyImageExtractor struct{}

func (sharedDummyImageExtractor) ExtractImageText(context.Context, backend.ImageExtractRequest) (string, backend.ImageExtractUsage, error) {
	return "shared", backend.ImageExtractUsage{}, nil
}

// parseTestDSNParts splits testDSN into connection coordinates for shared-DB
// registration (mirrors cloneTenantForID's parsing).
func parseTestDSNParts(t *testing.T) (host string, port int, user, pass, name string) {
	t.Helper()
	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)
		}
	}
	return host, port, parsed.User, parsed.Passwd, parsed.DBName
}

// Shared-schema tenants get the narrow image-extract enablement: the backend
// keeps only the resolver capability flag (no process-default extractor)
// while database auto-embedding, app semantic tasks, and audio extraction
// stay disabled exactly as before.
func TestPoolCreateBackendSharedTenantKeepsImageExtractCapability(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())
	initTenantPoolSchema(t, testDSN)
	// This focused pool fixture otherwise uses standalone tenant tables while
	// marking the placement shared. Backfill scheduling is the first scoped
	// semantic_tasks write in this test, so give that table its shared column.
	if _, err := metaStore.DB().Exec(`ALTER TABLE semantic_tasks ADD COLUMN fs_id BIGINT NOT NULL DEFAULT 0`); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	pool, tenantRecord := newTestPoolAndTenantWithConfig(t, PoolConfig{
		MaxTenants: 2,
		BackendOptions: backend.Options{
			DatabaseAutoEmbedding:   true,
			AppSemanticTasksEnabled: true,
			AsyncImageExtract: backend.AsyncImageExtractOptions{
				Enabled:    true,
				Capability: true,
				Extractor:  sharedDummyImageExtractor{},
			},
			AsyncAudioExtract: backend.AsyncAudioExtractOptions{Enabled: true, Extractor: &poolDummyAudioExtractor{}},
		},
	}, "tenant-shared-img-capability")
	pool.SetMetaStore(metaStore)
	resolveCalls := 0
	pool.SetEmbeddingRuntimeResolver(func(_ context.Context, tenantID string) (backend.AppEmbeddingRuntime, bool, error) {
		resolveCalls++
		if tenantID != tenantRecord.ID {
			t.Fatalf("embedding resolver tenant=%q, want %q", tenantID, tenantRecord.ID)
		}
		return backend.AppEmbeddingRuntime{Client: embedding.NopClient{}, Generation: 7, ConfigRevision: 9, TenantManaged: true}, true, nil
	})

	tenantRecord.Provider = ProviderTiDBCloudNativeShared
	now := time.Now().UTC()
	tenantRecord.CreatedAt = now
	tenantRecord.UpdatedAt = now
	if err := metaStore.InsertTenant(ctx, tenantRecord); err != nil {
		t.Fatal(err)
	}
	fsID, err := metaStore.EnsureFsID(ctx, tenantRecord.ID)
	if err != nil {
		t.Fatal(err)
	}
	dsnHost, dsnPort, dsnUser, dsnPass, dsnName := parseTestDSNParts(t)
	passwordCipher, err := pool.Encrypt(ctx, []byte(dsnPass))
	if err != nil {
		t.Fatal(err)
	}
	dbID, err := metaStore.RegisterSharedDB(ctx, &meta.SharedDB{
		TiDBCloudOrganizationID: "org-shared-img",
		Host:                    dsnHost,
		Port:                    dsnPort,
		User:                    dsnUser,
		PasswordCipher:          passwordCipher,
		Name:                    dsnName,
		MaxTenants:              100,
		Status:                  meta.SharedDBStatusActive,
	})
	if err != nil {
		t.Fatal(err)
	}
	// Mark the shared DB schema current so the handle open skips the schema
	// ensure (the test database keeps its standalone shape; this test never
	// executes scoped file queries).
	if _, err := metaStore.DB().ExecContext(ctx,
		`UPDATE db_pool SET schema_version = ? WHERE db_id = ?`,
		schema.CurrentSharedTiDBSchemaVersion, dbID); err != nil {
		t.Fatal(err)
	}
	if err := metaStore.UpsertTenantPlacement(ctx, &meta.TenantPlacement{
		FsID: fsID, DbID: dbID, Placement: meta.PlacementShared, SchemaShape: meta.SchemaShapeShared,
	}); err != nil {
		t.Fatal(err)
	}

	b, store, _, sharedLease, err := pool.createBackend(ctx, tenantRecord)
	if err != nil {
		t.Fatalf("createBackend(): %v", err)
	}
	t.Cleanup(func() {
		b.Close()
		_ = store.Close()
		if sharedLease != nil {
			sharedLease.Release()
		}
	})

	if !b.SupportsImageExtractTasks() {
		t.Fatal("shared backend must keep the image-extract capability")
	}
	if b.SupportsAsyncImageExtract() {
		t.Fatal("shared backend must not wire the process-default image extractor")
	}
	if b.SupportsAsyncAudioExtract() {
		t.Fatal("shared backend must keep audio extraction disabled")
	}
	if b.UsesDatabaseAutoEmbedding() {
		t.Fatal("shared backend must keep database auto-embedding disabled")
	}
	if !b.SupportsAppEmbeddingTasks() {
		t.Fatal("shared backend must enable app embedding with a resolved tenant runtime")
	}
	if resolveCalls != 1 {
		t.Fatalf("embedding resolver calls=%d, want 1", resolveCalls)
	}
	runtime, ok := b.AppEmbeddingRuntime()
	if !ok || runtime.Generation != 7 || runtime.ConfigRevision != 9 || !runtime.TenantManaged {
		t.Fatalf("app embedding runtime=%+v ok=%v, want generation=7 config_revision=9", runtime, ok)
	}
	var backfills int
	if err := store.DB().QueryRow(`SELECT COUNT(*) FROM semantic_tasks
		WHERE fs_id = ? AND task_type = ? AND resource_version = ?`, fsID, semantic.TaskTypeEmbedBackfill, 9).Scan(&backfills); err != nil {
		t.Fatal(err)
	}
	if backfills != 1 {
		t.Fatalf("shared embedding backfill tasks=%d, want 1", backfills)
	}
}