package tidbruntime

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"
)

func TestReportTiDBProvisionTimingWritesConfiguredFile(t *testing.T) {
	timingFile := filepath.Join(t.TempDir(), "timing.log")
	t.Setenv(tidbTestTimingEnv, "1")
	t.Setenv(tidbTestTimingFileEnv, timingFile)

	reportTiDBProvisionTiming(t, "clone_database", 12*time.Millisecond, nil)

	contents, err := os.ReadFile(timingFile)
	require.NoError(t, err)
	require.Contains(t, string(contents),
		"test=\"TestReportTiDBProvisionTimingWritesConfiguredFile\" operation=clone_database status=passed duration_ms=12",
	)
}

func TestStartIsolatedTiDBIgnoresConfiguredTemplateDatabase(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEST_DSN", "root@tcp(127.0.0.1:4000)/root")
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "agentteam_tidb_template_configured")

	originalStartDatabase := startTiDBDatabaseFn
	t.Cleanup(func() {
		startTiDBDatabaseFn = originalStartDatabase
	})

	startTiDBDatabaseFn = func(t *testing.T, rootDSN string, templateDatabase string) string {
		t.Helper()
		require.Equal(t, "root@tcp(127.0.0.1:4000)/root", rootDSN)
		require.Empty(t, templateDatabase)
		return "root@tcp(127.0.0.1:4000)/isolated"
	}

	require.Equal(t, "root@tcp(127.0.0.1:4000)/isolated", StartIsolatedTiDB(t))
}

func TestInitializedTiDBTemplateDatabaseUsesConfiguredTemplateWhenReady(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "agentteam_tidb_template_configured")

	restore := stubInitializedTemplateDatabaseDependencies(t)
	defer restore()

	rootDSN := "root@tcp(127.0.0.1:4000)/root"
	tidbTemplateReadyOnRootDSNFn = func(ctx context.Context, gotRootDSN string, templateDatabase string) (bool, error) {
		require.Equal(t, rootDSN, gotRootDSN)
		require.Equal(t, "agentteam_tidb_template_configured", templateDatabase)
		return true, nil
	}
	prepareMigratedTiDBTemplateFn = func(ctx context.Context, gotRootDSN string) (string, error) {
		t.Fatal("PrepareMigratedTiDBTemplate should not run when the configured template is ready")
		return "", nil
	}

	templateDatabase, err := initializedTiDBTemplateDatabase(context.Background(), rootDSN)
	require.NoError(t, err)
	require.Equal(t, "agentteam_tidb_template_configured", templateDatabase)
}

func TestInitializedTiDBTemplateDatabaseFallsBackWhenConfiguredTemplateNotReady(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "agentteam_tidb_template_configured")

	restore := stubInitializedTemplateDatabaseDependencies(t)
	defer restore()

	rootDSN := "root@tcp(127.0.0.1:4000)/root"
	tidbTemplateReadyOnRootDSNFn = func(ctx context.Context, gotRootDSN string, templateDatabase string) (bool, error) {
		require.Equal(t, rootDSN, gotRootDSN)
		require.Equal(t, "agentteam_tidb_template_configured", templateDatabase)
		return false, nil
	}
	prepareMigratedTiDBTemplateFn = func(ctx context.Context, gotRootDSN string) (string, error) {
		require.Equal(t, rootDSN, gotRootDSN)
		return "agentteam_tidb_template_prepared", nil
	}

	templateDatabase, err := initializedTiDBTemplateDatabase(context.Background(), rootDSN)
	require.NoError(t, err)
	require.Equal(t, "agentteam_tidb_template_prepared", templateDatabase)
}

func TestTiDBTemplateReadyOnRootDSNRequiresFingerprintMatch(t *testing.T) {
	restore := stubConfiguredTemplateReadyDependencies(t)
	defer restore()

	events := make([]string, 0)
	originalOpenSQLDB := openSQLDBFn
	openSQLDBFn = func(driverName string, dataSourceName string) (*sql.DB, error) {
		events = append(events, "open")
		require.Equal(t, "mysql", driverName)
		require.Contains(t, dataSourceName, "root@tcp(127.0.0.1:4000)/")
		db, mock, err := sqlmock.New(sqlmock.MonitorPingsOption(true))
		require.NoError(t, err)
		mock.ExpectPing()
		mock.ExpectClose()
		t.Cleanup(func() {
			require.NoError(t, mock.ExpectationsWereMet())
		})
		return db, nil
	}
	defer func() {
		openSQLDBFn = originalOpenSQLDB
	}()

	rootDSN := "root@tcp(127.0.0.1:4000)/root"
	tidbDatabaseExistsFn = func(ctx context.Context, adminDB *sql.DB, databaseName string) (bool, error) {
		events = append(events, "exists")
		require.NotNil(t, adminDB)
		require.Equal(t, "agentteam_tidb_template_configured", databaseName)
		return true, nil
	}
	tidbDatabaseSchemaCurrentFn = func(ctx context.Context, gotRootDSN string, databaseName string) (bool, error) {
		events = append(events, "schema-current")
		require.Equal(t, rootDSN, gotRootDSN)
		require.Equal(t, "agentteam_tidb_template_configured", databaseName)
		return true, nil
	}
	tidbTemplateSourceFingerprintFn = func() (string, error) {
		events = append(events, "source-fingerprint")
		return "expected-fingerprint", nil
	}
	tidbTemplateFingerprintMatchesFn = func(ctx context.Context, adminDB *sql.DB, templateDatabase string, want string) (bool, error) {
		events = append(events, "fingerprint")
		require.NotNil(t, adminDB)
		require.Equal(t, "agentteam_tidb_template_configured", templateDatabase)
		require.Equal(t, "expected-fingerprint", want)
		return false, nil
	}

	ready, err := tidbTemplateReadyOnRootDSN(context.Background(), rootDSN, "agentteam_tidb_template_configured")
	require.NoError(t, err)
	require.False(t, ready)
	require.Equal(t, []string{"open", "exists", "schema-current", "source-fingerprint", "fingerprint"}, events)
}

func TestCloneTiDBDatabaseFromDefaultTemplateUsesCloneSlotWithoutTemplateLock(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "")

	restore := stubCloneTiDBDatabaseFromTemplateDependencies(t)
	defer restore()
	events := make([]string, 0)
	restoreOpen := stubOpenSQLDBWithPing(t, &events)
	defer restoreOpen()

	withSharedTiDBCloneSlotContextFn = func(ctx context.Context, timeout time.Duration, fn func() error) error {
		events = append(events, "clone-slot-enter")
		err := fn()
		events = append(events, "clone-slot-exit")
		return err
	}
	cloneTiDBDatabaseWithCacheFn = func(ctx context.Context, adminDB *sql.DB, cloneAdminDSN string, cacheNamespace string, templateDatabase string, databaseName string) error {
		events = append(events, "clone")
		require.Equal(t, "root@tcp(127.0.0.1:4000)/", cloneAdminDSN)
		require.Equal(t, "root@tcp(127.0.0.1:4000)/", cacheNamespace)
		require.Equal(t, defaultTiDBTemplateDatabase, templateDatabase)
		require.Equal(t, "agentteam_tidb_test_target", databaseName)
		return nil
	}

	err := cloneTiDBDatabaseFromTemplateWithFreshAdmin(context.Background(), "root@tcp(127.0.0.1:4000)/", defaultTiDBTemplateDatabase, "agentteam_tidb_test_target", time.Second)
	require.NoError(t, err)
	require.Equal(t, []string{
		"clone-slot-enter",
		"open",
		"clone",
		"clone-slot-exit",
	}, events)
}

func TestCloneTiDBDatabaseFromConfiguredTemplateDoesNotTakeDefaultTemplateLock(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "agentteam_tidb_template_configured")

	restore := stubCloneTiDBDatabaseFromTemplateDependencies(t)
	defer restore()
	events := make([]string, 0)
	restoreOpen := stubOpenSQLDBWithPing(t, &events)
	defer restoreOpen()

	withSharedTiDBCloneSlotContextFn = func(ctx context.Context, timeout time.Duration, fn func() error) error {
		events = append(events, "clone-slot")
		return fn()
	}
	cloneTiDBDatabaseWithCacheFn = func(ctx context.Context, adminDB *sql.DB, cloneAdminDSN string, cacheNamespace string, templateDatabase string, databaseName string) error {
		events = append(events, "clone")
		require.Empty(t, cacheNamespace)
		require.Equal(t, "agentteam_tidb_template_configured", templateDatabase)
		return nil
	}

	err := cloneTiDBDatabaseFromTemplateWithFreshAdmin(context.Background(), "root@tcp(127.0.0.1:4000)/", "agentteam_tidb_template_configured", "agentteam_tidb_test_target", time.Second)
	require.NoError(t, err)
	require.Equal(t, []string{"clone-slot", "open", "clone"}, events)
}

func TestCloneTiDBDatabaseFromTemplateWithFreshAdminRetriesTransientConnectionError(t *testing.T) {
	t.Setenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE", "")

	restore := stubCloneTiDBDatabaseFromTemplateDependencies(t)
	defer restore()
	originalDropDatabase := dropTiDBDatabaseFn
	defer func() {
		dropTiDBDatabaseFn = originalDropDatabase
	}()

	events := make([]string, 0)
	restoreOpen := stubOpenSQLDBWithPing(t, &events)
	defer restoreOpen()
	withSharedTiDBCloneSlotContextFn = func(ctx context.Context, timeout time.Duration, fn func() error) error {
		events = append(events, "clone-slot-enter")
		err := fn()
		events = append(events, "clone-slot-exit")
		return err
	}
	cloneAttempts := 0
	cloneTiDBDatabaseWithCacheFn = func(ctx context.Context, adminDB *sql.DB, cloneAdminDSN string, cacheNamespace string, templateDatabase string, databaseName string) error {
		cloneAttempts++
		events = append(events, fmt.Sprintf("clone-%d", cloneAttempts))
		require.NotNil(t, adminDB)
		require.Equal(t, "root@tcp(127.0.0.1:4000)/", cloneAdminDSN)
		require.Equal(t, "root@tcp(127.0.0.1:4000)/", cacheNamespace)
		require.Equal(t, defaultTiDBTemplateDatabase, templateDatabase)
		require.Equal(t, "agentteam_tidb_test_target", databaseName)
		if cloneAttempts == 1 {
			return fmt.Errorf("clone tidb table structure template.messages: %w", mysql.ErrInvalidConn)
		}
		return nil
	}
	dropTiDBDatabaseFn = func(ctx context.Context, adminDB *sql.DB, databaseName string) error {
		events = append(events, "drop")
		require.NotNil(t, adminDB)
		require.Equal(t, "agentteam_tidb_test_target", databaseName)
		return nil
	}

	err := cloneTiDBDatabaseFromTemplateWithFreshAdmin(
		context.Background(),
		"root@tcp(127.0.0.1:4000)/",
		defaultTiDBTemplateDatabase,
		"agentteam_tidb_test_target",
		time.Second,
	)
	require.NoError(t, err)
	require.Equal(t, []string{
		"clone-slot-enter",
		"open",
		"clone-1",
		"open",
		"drop",
		"open",
		"clone-2",
		"clone-slot-exit",
	}, events)
}

func TestRetryTiDBProvisioningRetriesTransientConnectionError(t *testing.T) {
	attempts := 0
	cleanups := 0

	err := retryTiDBProvisioning(context.Background(), "agentteam_tidb_test_target", func(ctx context.Context) error {
		attempts++
		if attempts == 1 {
			return fmt.Errorf("clone tidb table structure template.messages: %w", mysql.ErrInvalidConn)
		}
		return nil
	}, func(ctx context.Context) error {
		cleanups++
		return nil
	})

	require.NoError(t, err)
	require.Equal(t, 2, attempts)
	require.Equal(t, 1, cleanups)
}

func TestRetryTiDBProvisioningRetriesTiDBSchemaLeaseError(t *testing.T) {
	attempts := 0
	cleanups := 0

	err := retryTiDBProvisioning(context.Background(), "agentteam_tidb_test_target", func(ctx context.Context) error {
		attempts++
		if attempts == 1 {
			return fmt.Errorf("create tidb database agentteam_tidb_test_target: %w", &mysql.MySQLError{
				Number:  8027,
				Message: "Information schema is out of date: schema failed to update in 1 lease, please make sure TiDB can connect to TiKV",
			})
		}
		return nil
	}, func(ctx context.Context) error {
		cleanups++
		return nil
	})

	require.NoError(t, err)
	require.Equal(t, 2, attempts)
	require.Equal(t, 1, cleanups)
}

func TestRetryTiDBProvisioningRetriesSetupDeadline(t *testing.T) {
	attempts := 0
	cleanups := 0

	err := retryTiDBProvisioning(context.Background(), "agentteam_tidb_test_target", func(ctx context.Context) error {
		attempts++
		if attempts == 1 {
			return fmt.Errorf("clone tidb table structure template.messages: %w", context.DeadlineExceeded)
		}
		return nil
	}, func(ctx context.Context) error {
		cleanups++
		return nil
	})

	require.NoError(t, err)
	require.Equal(t, 2, attempts)
	require.Equal(t, 1, cleanups)
}

func TestRetryTiDBProvisioningDoesNotRetrySemanticError(t *testing.T) {
	attempts := 0
	cleanups := 0
	wantErr := fmt.Errorf("clone tidb table structure template.messages: syntax error")

	err := retryTiDBProvisioning(context.Background(), "agentteam_tidb_test_target", func(ctx context.Context) error {
		attempts++
		return wantErr
	}, func(ctx context.Context) error {
		cleanups++
		return nil
	})

	require.ErrorIs(t, err, wantErr)
	require.Equal(t, 1, attempts)
	require.Equal(t, 0, cleanups)
}

func TestDropTiDBDatabaseWithAdminDSNRetriesTransientConnectionError(t *testing.T) {
	originalDropDatabase := dropTiDBDatabaseFn
	defer func() {
		dropTiDBDatabaseFn = originalDropDatabase
	}()

	events := make([]string, 0)
	restoreOpen := stubOpenSQLDBWithPing(t, &events)
	defer restoreOpen()

	dropAttempts := 0
	dropTiDBDatabaseFn = func(ctx context.Context, adminDB *sql.DB, databaseName string) error {
		dropAttempts++
		events = append(events, fmt.Sprintf("drop-%d", dropAttempts))
		require.NotNil(t, adminDB)
		require.Equal(t, "agentteam_tidb_test_target", databaseName)
		if dropAttempts == 1 {
			return fmt.Errorf("drop tidb database agentteam_tidb_test_target: %w", mysql.ErrInvalidConn)
		}
		return nil
	}

	err := dropTiDBDatabaseWithAdminDSN(context.Background(), "root@tcp(127.0.0.1:4000)/", "agentteam_tidb_test_target")
	require.NoError(t, err)
	require.Equal(t, []string{
		"open",
		"drop-1",
		"open",
		"drop-2",
	}, events)
}

func stubInitializedTemplateDatabaseDependencies(t *testing.T) func() {
	t.Helper()

	originalReady := tidbTemplateReadyOnRootDSNFn
	originalPrepare := prepareMigratedTiDBTemplateFn
	resetPreparedTiDBTemplateCache()

	tidbTemplateReadyOnRootDSNFn = func(ctx context.Context, rootDSN string, templateDatabase string) (bool, error) {
		t.Fatal("tidbTemplateReadyOnRootDSNFn should not be called")
		return false, nil
	}
	prepareMigratedTiDBTemplateFn = func(ctx context.Context, rootDSN string) (string, error) {
		t.Fatal("prepareMigratedTiDBTemplateFn should not be called")
		return "", nil
	}

	return func() {
		tidbTemplateReadyOnRootDSNFn = originalReady
		prepareMigratedTiDBTemplateFn = originalPrepare
		resetPreparedTiDBTemplateCache()
	}
}

func stubConfiguredTemplateReadyDependencies(t *testing.T) func() {
	t.Helper()

	originalDatabaseExists := tidbDatabaseExistsFn
	originalSchemaCurrent := tidbDatabaseSchemaCurrentFn
	originalSourceFingerprint := tidbTemplateSourceFingerprintFn
	originalFingerprintMatches := tidbTemplateFingerprintMatchesFn

	tidbDatabaseExistsFn = func(ctx context.Context, adminDB *sql.DB, databaseName string) (bool, error) {
		t.Fatal("tidbDatabaseExistsFn should not be called")
		return false, nil
	}
	tidbDatabaseSchemaCurrentFn = func(ctx context.Context, rootDSN string, databaseName string) (bool, error) {
		t.Fatal("tidbDatabaseSchemaCurrentFn should not be called")
		return false, nil
	}
	tidbTemplateSourceFingerprintFn = func() (string, error) {
		t.Fatal("tidbTemplateSourceFingerprintFn should not be called")
		return "", nil
	}
	tidbTemplateFingerprintMatchesFn = func(ctx context.Context, adminDB *sql.DB, templateDatabase string, want string) (bool, error) {
		t.Fatal("tidbTemplateFingerprintMatchesFn should not be called")
		return false, nil
	}

	return func() {
		tidbDatabaseExistsFn = originalDatabaseExists
		tidbDatabaseSchemaCurrentFn = originalSchemaCurrent
		tidbTemplateSourceFingerprintFn = originalSourceFingerprint
		tidbTemplateFingerprintMatchesFn = originalFingerprintMatches
	}
}

func stubCloneTiDBDatabaseFromTemplateDependencies(t *testing.T) func() {
	t.Helper()

	originalCloneSlot := withSharedTiDBCloneSlotContextFn
	originalClone := cloneTiDBDatabaseWithCacheFn

	withSharedTiDBCloneSlotContextFn = func(ctx context.Context, timeout time.Duration, fn func() error) error {
		t.Fatal("withSharedTiDBCloneSlotContextFn should not be called")
		return nil
	}
	cloneTiDBDatabaseWithCacheFn = func(ctx context.Context, adminDB *sql.DB, cloneAdminDSN string, cacheNamespace string, templateDatabase string, databaseName string) error {
		t.Fatal("cloneTiDBDatabaseWithCacheFn should not be called")
		return nil
	}

	return func() {
		withSharedTiDBCloneSlotContextFn = originalCloneSlot
		cloneTiDBDatabaseWithCacheFn = originalClone
	}
}

func stubOpenSQLDBWithPing(t *testing.T, events *[]string) func() {
	t.Helper()

	originalOpenSQLDB := openSQLDBFn
	mocks := make([]sqlmock.Sqlmock, 0)
	openSQLDBFn = func(driverName string, dataSourceName string) (*sql.DB, error) {
		if events != nil {
			*events = append(*events, "open")
		}
		require.Equal(t, "mysql", driverName)
		require.Equal(t, "root@tcp(127.0.0.1:4000)/", dataSourceName)
		db, mock, err := sqlmock.New(sqlmock.MonitorPingsOption(true))
		require.NoError(t, err)
		mock.ExpectPing()
		mock.ExpectClose()
		mocks = append(mocks, mock)
		return db, nil
	}

	return func() {
		openSQLDBFn = originalOpenSQLDB
		for _, mock := range mocks {
			require.NoError(t, mock.ExpectationsWereMet())
		}
	}
}