package meta

import (
	"context"
	"database/sql"
	"errors"
	"sync"
	"testing"
	"time"
)

func TestCountTiDBCloudOrganizationTenantsCountsAllOwnedTenants(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()
	tests := []struct {
		id         string
		provider   string
		status     TenantStatus
		org        string
		poolStatus TenantPoolBindingStatus
		want       bool
	}{
		{id: "native-pending", provider: tidbCloudNativeProvider, status: TenantPending, org: "org-a", want: true},
		{id: "native-active", provider: tidbCloudNativeProvider, status: TenantActive, org: "org-a", want: true},
		{id: "shared-failed", provider: tidbCloudNativeSharedProvider, status: TenantFailed, org: "org-a", want: true},
		{id: "shared-deleting", provider: tidbCloudNativeSharedProvider, status: TenantDeleting, org: "org-a", want: true},
		{id: "native-paid", provider: tidbCloudNativeProvider, status: TenantActive, org: "org-a", want: true},
		{id: "native-deleted", provider: tidbCloudNativeProvider, status: TenantDeleted, org: "org-a"},
		{id: "shared-other-org", provider: tidbCloudNativeSharedProvider, status: TenantActive, org: "org-b"},
		{id: "native-free-pool-inventory", provider: tidbCloudNativeProvider, status: TenantActive, org: "org-a", poolStatus: TenantPoolBindingFree},
		{id: "shared-free-pool-inventory", provider: tidbCloudNativeSharedProvider, status: TenantActive, org: "org-a", poolStatus: TenantPoolBindingFree},
	}
	want := 0
	for _, tt := range tests {
		insertFreeQuotaTestTenant(t, s, tt.id, tt.status, tt.provider)
		switch tt.provider {
		case tidbCloudNativeProvider:
			binding := &TenantTiDBCloudOrgBinding{
				TenantID: tt.id, OrganizationID: tt.org, ClusterID: "cluster-" + tt.id,
			}
			if tt.poolStatus != "" {
				binding.PoolID = "pool-" + tt.id
				binding.PoolStatus = tt.poolStatus
			}
			if err := s.UpsertTenantTiDBCloudOrgBinding(ctx, binding); err != nil {
				t.Fatal(err)
			}
		case tidbCloudNativeSharedProvider:
			insertSharedTenantPlacementForExistingTenantTest(t, s, tt.id, tt.org)
			if tt.poolStatus != "" {
				if err := s.UpsertTenantPoolMembership(ctx, &TenantPoolMembership{
					TenantID: tt.id, PoolID: "pool-" + tt.id,
					TiDBCloudOrganizationID: tt.org, PoolStatus: tt.poolStatus,
				}); err != nil {
					t.Fatal(err)
				}
			}
		default:
			t.Fatalf("unsupported provider %q", tt.provider)
		}
		if tt.want {
			want++
		}
	}
	got, err := s.CountTiDBCloudOrganizationTenants(ctx, "org-a")
	if err != nil {
		t.Fatal(err)
	}
	if got != want {
		t.Fatalf("organization tenant count = %d, want %d", got, want)
	}
}

func TestMigrateDropsObsoleteTenantQuotaColumnAndKeepsLegacyCompatibilityColumns(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()
	ensureColumn := func(name, definition string) {
		t.Helper()
		var count int
		if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM information_schema.columns
			WHERE table_schema = DATABASE() AND table_name = 'tenant_quota_config' AND column_name = ?`, name).Scan(&count); err != nil {
			t.Fatal(err)
		}
		if count > 0 {
			return
		}
		if _, err := s.DB().ExecContext(ctx, `ALTER TABLE tenant_quota_config ADD COLUMN `+name+` `+definition); err != nil {
			t.Fatal(err)
		}
	}
	ensureColumn("quota_limits_overridden", "TINYINT(1) NOT NULL DEFAULT 1")
	ensureColumn("tidbcloud_spending_limit", "BIGINT NULL")
	ensureColumn("tidbcloud_spending_limit_checked_at", "DATETIME(3) NULL")

	if err := s.migrate(); err != nil {
		t.Fatal(err)
	}
	if err := s.migrate(); err != nil {
		t.Fatal(err)
	}
	for name, want := range map[string]int{
		"quota_limits_overridden":             0,
		"tidbcloud_spending_limit":            1,
		"tidbcloud_spending_limit_checked_at": 1,
	} {
		var got int
		if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM information_schema.columns
			WHERE table_schema = DATABASE() AND table_name = 'tenant_quota_config' AND column_name = ?`, name).Scan(&got); err != nil {
			t.Fatal(err)
		}
		if got != want {
			t.Fatalf("tenant_quota_config.%s exists = %d, want %d", name, got, want)
		}
	}
}

func TestMigrateUpdatesTenantQuotaLLMDefaultsWithoutChangingExistingRows(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()

	if _, err := s.DB().ExecContext(ctx, `INSERT INTO tenant_quota_config
		(tenant_id, max_media_llm_files, max_video_llm_files) VALUES ('existing-quota', 700, 40)`); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `ALTER TABLE tenant_quota_config
		MODIFY COLUMN max_media_llm_files BIGINT NOT NULL DEFAULT 1000,
		MODIFY COLUMN max_video_llm_files BIGINT NOT NULL DEFAULT 50`); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		_, _ = s.DB().ExecContext(context.Background(), `ALTER TABLE tenant_quota_config
			MODIFY COLUMN max_media_llm_files BIGINT NOT NULL DEFAULT 100,
			MODIFY COLUMN max_video_llm_files BIGINT NOT NULL DEFAULT 30`)
	})

	if err := s.migrate(); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `INSERT INTO tenant_quota_config (tenant_id) VALUES ('new-quota')`); err != nil {
		t.Fatal(err)
	}
	for tenantID, want := range map[string][2]int64{
		"existing-quota": {700, 40},
		"new-quota":      {100, 30},
	} {
		var media, video int64
		if err := s.DB().QueryRowContext(ctx, `SELECT max_media_llm_files, max_video_llm_files
			FROM tenant_quota_config WHERE tenant_id = ?`, tenantID).Scan(&media, &video); err != nil {
			t.Fatal(err)
		}
		if media != want[0] || video != want[1] {
			t.Fatalf("%s LLM quota = %d/%d, want %d/%d", tenantID, media, video, want[0], want[1])
		}
	}
}

func TestLegacyFreeTenantClassificationRollbackCompatibility(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()

	// The new counter uses durable organization ownership and does not depend on
	// the legacy spending-limit classification written by an old binary.
	const oldWriterTenantID = "legacy-writer-free"
	insertFreeQuotaTestTenant(t, s, oldWriterTenantID, TenantActive, tidbCloudNativeProvider)
	if err := s.UpsertTenantTiDBCloudOrgBinding(ctx, &TenantTiDBCloudOrgBinding{
		TenantID: oldWriterTenantID, OrganizationID: "org-legacy", ClusterID: "cluster-legacy",
	}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `INSERT INTO tenant_quota_config
		(tenant_id, tidbcloud_spending_limit) VALUES (?, 0)`, oldWriterTenantID); err != nil {
		t.Fatal(err)
	}
	count, err := s.CountTiDBCloudOrganizationTenants(ctx, "org-legacy")
	if err != nil {
		t.Fatal(err)
	}
	if count != 1 {
		t.Fatalf("organization tenant count = %d, want 1", count)
	}

	// New writers shadow the free classification into the retained column so a
	// rollback within the compatibility window preserves the legacy behavior.
	storage, fileSize, fileCount := int64(3<<30), int64(300<<20), int64(1000)
	const newWriterTenantID = "rollback-shadow-free"
	if err := s.InsertTiDBCloudFreeTenant(ctx, newFreeQuotaTestTenant(newWriterTenantID), QuotaConfigPatch{
		MaxStorageBytes: &storage, MaxFileSizeBytes: &fileSize, MaxFileCount: &fileCount,
	}); err != nil {
		t.Fatal(err)
	}
	var legacySpending sql.NullInt64
	if err := s.DB().QueryRowContext(ctx, `SELECT tidbcloud_spending_limit FROM tenant_quota_config
		WHERE tenant_id = ?`, newWriterTenantID).Scan(&legacySpending); err != nil {
		t.Fatal(err)
	}
	if !legacySpending.Valid || legacySpending.Int64 != 0 {
		t.Fatalf("new-writer legacy classification = %+v, want 0", legacySpending)
	}

	tx, err := s.DB().BeginTx(ctx, nil)
	if err != nil {
		t.Fatal(err)
	}
	if err := setLegacyTiDBCloudFreeClassificationTx(ctx, tx, newWriterTenantID, false); err != nil {
		_ = tx.Rollback()
		t.Fatal(err)
	}
	if err := tx.Commit(); err != nil {
		t.Fatal(err)
	}
	if err := s.DB().QueryRowContext(ctx, `SELECT tidbcloud_spending_limit FROM tenant_quota_config
		WHERE tenant_id = ?`, newWriterTenantID).Scan(&legacySpending); err != nil {
		t.Fatal(err)
	}
	if legacySpending.Valid {
		t.Fatalf("paid claim legacy classification = %+v, want NULL", legacySpending)
	}
}

func insertSharedTenantPlacementForExistingTenantTest(t *testing.T, s *Store, tenantID, organizationID string) {
	t.Helper()
	ctx := context.Background()
	fsID, err := s.EnsureFsID(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	dbID, err := s.RegisterSharedDB(ctx, &SharedDB{
		TiDBCloudOrganizationID: organizationID,
		Host:                    tenantID + ".shared.example.com",
		Port:                    4000,
		User:                    "root",
		PasswordCipher:          []byte("cipher"),
		Name:                    "shared_" + tenantID,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.UpsertTenantPlacement(ctx, &TenantPlacement{
		FsID: fsID, DbID: dbID, Placement: PlacementShared, SchemaShape: SchemaShapeShared,
	}); err != nil {
		t.Fatal(err)
	}
}

func TestWithTiDBCloudFreeQuotaLockSerializesAndReturnsBusy(t *testing.T) {
	s := newControlStore(t)
	originalTimeout := tidbCloudFreeQuotaLockTimeoutSeconds
	tidbCloudFreeQuotaLockTimeoutSeconds = 0
	t.Cleanup(func() { tidbCloudFreeQuotaLockTimeoutSeconds = originalTimeout })

	acquired := make(chan struct{})
	release := make(chan struct{})
	var firstErr error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		firstErr = s.WithTiDBCloudFreeQuotaLock(context.Background(), "org-lock", func(context.Context) error {
			close(acquired)
			<-release
			return nil
		})
	}()
	<-acquired
	err := s.WithTiDBCloudFreeQuotaLock(context.Background(), "org-lock", func(context.Context) error {
		t.Fatal("busy lock callback executed")
		return nil
	})
	close(release)
	wg.Wait()
	if firstErr != nil {
		t.Fatalf("first lock: %v", firstErr)
	}
	if !errors.Is(err, ErrTiDBCloudFreeQuotaBusy) {
		t.Fatalf("second lock error = %v, want busy", err)
	}
}

func TestInsertTiDBCloudFreeTenantPersistsExplicitQuota(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()
	if _, err := s.DB().ExecContext(ctx, `ALTER TABLE tenant_quota_config
		MODIFY COLUMN max_media_llm_files BIGINT NOT NULL DEFAULT 1000,
		MODIFY COLUMN max_video_llm_files BIGINT NOT NULL DEFAULT 50`); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		_, _ = s.DB().ExecContext(context.Background(), `ALTER TABLE tenant_quota_config
			MODIFY COLUMN max_media_llm_files BIGINT NOT NULL DEFAULT 100,
			MODIFY COLUMN max_video_llm_files BIGINT NOT NULL DEFAULT 30`)
	})
	storage, fileSize, fileCount := int64(3<<30), int64(300<<20), int64(1000)
	quota := QuotaConfigPatch{
		MaxStorageBytes:  &storage,
		MaxFileSizeBytes: &fileSize,
		MaxFileCount:     &fileCount,
	}
	const tenantID = "free-tenant"
	if err := s.InsertTiDBCloudFreeTenant(ctx, newFreeQuotaTestTenant(tenantID), quota); err != nil {
		t.Fatal(err)
	}
	cfg, err := s.GetQuotaConfig(ctx, tenantID)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.MaxStorageBytes != storage || cfg.MaxFileSizeBytes != fileSize || cfg.MaxFileCount != fileCount ||
		cfg.MaxMediaLLMFiles != 100 || cfg.MaxVideoLLMFiles != 30 {
		t.Fatalf("quota = %+v", cfg)
	}
	var legacySpending sql.NullInt64
	if err := s.DB().QueryRowContext(ctx,
		`SELECT tidbcloud_spending_limit FROM tenant_quota_config WHERE tenant_id = ?`, tenantID).Scan(&legacySpending); err != nil {
		t.Fatal(err)
	}
	if !legacySpending.Valid || legacySpending.Int64 != 0 {
		t.Fatalf("legacy free classification = %+v, want 0", legacySpending)
	}
}

func TestInsertTiDBCloudFreeTenantRollsBackEveryRow(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()
	const tenantID = "free-tenant-rollback"
	if _, err := s.DB().ExecContext(ctx, `INSERT INTO tenant_quota_config (tenant_id) VALUES (?)`, tenantID); err != nil {
		t.Fatal(err)
	}
	storage, fileSize, fileCount := int64(3<<30), int64(300<<20), int64(1000)
	err := s.InsertTiDBCloudFreeTenant(ctx, newFreeQuotaTestTenant(tenantID), QuotaConfigPatch{
		MaxStorageBytes:  &storage,
		MaxFileSizeBytes: &fileSize,
		MaxFileCount:     &fileCount,
	})
	if err == nil {
		t.Fatal("free tenant insert with duplicate quota row succeeded")
	}
	if _, err := s.GetTenant(ctx, tenantID); !errors.Is(err, ErrNotFound) {
		t.Fatalf("tenant lookup after rollback = %v, want not found", err)
	}
	var quotaRows int
	if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM tenant_quota_config WHERE tenant_id = ?`, tenantID).Scan(&quotaRows); err != nil {
		t.Fatal(err)
	}
	if quotaRows != 1 {
		t.Fatalf("quota rows after rollback = %d, want preexisting row only", quotaRows)
	}
}

func TestHasTenantPoolOwnershipUsesExplicitOwnershipMetadata(t *testing.T) {
	s := newControlStore(t)
	ctx := context.Background()

	insertFreeQuotaTestTenant(t, s, "ownership-none", TenantPending, tidbCloudNativeProvider)
	owned, err := s.HasTenantPoolOwnership(ctx, "ownership-none")
	if err != nil || owned {
		t.Fatalf("ownership-none = %v, err=%v, want false", owned, err)
	}

	insertFreeQuotaTestTenant(t, s, "ownership-native-binding-only", TenantPending, tidbCloudNativeProvider)
	if err := s.UpsertTenantTiDBCloudOrgBinding(ctx, &TenantTiDBCloudOrgBinding{
		TenantID: "ownership-native-binding-only", OrganizationID: "org-binding-only", ClusterID: "cluster-binding-only",
	}); err != nil {
		t.Fatal(err)
	}
	owned, err = s.HasTenantPoolOwnership(ctx, "ownership-native-binding-only")
	if err != nil || owned {
		t.Fatalf("native binding without pool = %v, err=%v, want false", owned, err)
	}

	insertFreeQuotaTestTenant(t, s, "ownership-native-pool", TenantPending, tidbCloudNativeProvider)
	if err := s.UpsertTenantTiDBCloudOrgBinding(ctx, &TenantTiDBCloudOrgBinding{
		TenantID: "ownership-native-pool", OrganizationID: "org-native-pool", ClusterID: "cluster-native-pool", PoolID: "pool-native",
	}); err != nil {
		t.Fatal(err)
	}
	owned, err = s.HasTenantPoolOwnership(ctx, "ownership-native-pool")
	if err != nil || !owned {
		t.Fatalf("native pool ownership = %v, err=%v, want true", owned, err)
	}

	insertFreeQuotaTestTenant(t, s, "ownership-membership", TenantPending, tidbCloudNativeSharedProvider)
	if _, err := s.DB().ExecContext(ctx, `INSERT INTO tenant_pool_memberships
		(tenant_id, pool_id, pool_status) VALUES ('ownership-membership', 'pool-membership', 'free')`); err != nil {
		t.Fatal(err)
	}
	owned, err = s.HasTenantPoolOwnership(ctx, "ownership-membership")
	if err != nil || !owned {
		t.Fatalf("membership ownership = %v, err=%v, want true", owned, err)
	}

	insertFreeQuotaTestTenant(t, s, "ownership-placement", TenantPending, tidbCloudNativeSharedProvider)
	fsID, err := s.EnsureFsID(ctx, "ownership-placement")
	if err != nil {
		t.Fatal(err)
	}
	dbID, err := s.RegisterSharedDB(ctx, &SharedDB{
		TiDBCloudOrganizationID: "org-ownership-placement", Host: "ownership-placement.example", Port: 4000,
		User: "root", PasswordCipher: []byte("cipher"), Name: "ownership_placement",
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := s.UpsertTenantPlacement(ctx, &TenantPlacement{
		FsID: fsID, DbID: dbID, Placement: PlacementShared, SchemaShape: SchemaShapeShared,
	}); err != nil {
		t.Fatal(err)
	}
	owned, err = s.HasTenantPoolOwnership(ctx, "ownership-placement")
	if err != nil || !owned {
		t.Fatalf("placement ownership = %v, err=%v, want true", owned, err)
	}
}

func insertFreeQuotaTestTenant(t *testing.T, s *Store, tenantID string, status TenantStatus, provider string) {
	t.Helper()
	now := time.Now().UTC()
	if err := s.InsertTenant(context.Background(), &Tenant{
		ID: tenantID, Status: status, Kind: TenantKindLive, Provider: provider,
		DBPasswordCipher: []byte{}, DBTLS: true, SchemaVersion: 1,
		CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatal(err)
	}
}

func newFreeQuotaTestTenant(tenantID string) *Tenant {
	now := time.Now().UTC()
	return &Tenant{
		ID: tenantID, Status: TenantPending, Kind: TenantKindLive, Provider: tidbCloudNativeProvider,
		DBPasswordCipher: []byte{}, DBTLS: true, SchemaVersion: 1,
		CreatedAt: now, UpdatedAt: now,
	}
}