package tidbruntime

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

var tidbTestDatabaseCounter uint64
var tidbTemplateMu sync.Mutex
var tidbPreparedTemplateRootDSN string
var tidbPreparedTemplateDatabase string

var (
	startTiDBDatabaseFn              = startTiDBDatabase
	tidbTemplateReadyOnRootDSNFn     = tidbTemplateReadyOnRootDSN
	prepareMigratedTiDBTemplateFn    = PrepareMigratedTiDBTemplate
	withSharedTiDBCloneSlotContextFn = withSharedTiDBCloneSlotContext
	cloneTiDBDatabaseWithCacheFn     = cloneTiDBDatabaseWithCache
)

const tidbTestDatabaseSetupTimeout = 2 * time.Minute
const tidbTemplatePrepareTimeout = 10 * time.Minute

const tidbProvisioningRetryAttempts = 3
const tidbProvisioningRetryDelay = 500 * time.Millisecond
const tidbInformationSchemaOutOfDateErrorNumber = 8027

const tidbTestTimingEnv = "AGENTTEAM_TIDB_TEST_TIMING"
const tidbTestTimingFileEnv = "AGENTTEAM_TIDB_TEST_TIMING_FILE"

// StartIsolatedTiDB returns a fresh TiDB database DSN created from
// AGENTTEAM_TIDB_TEST_DSN. When the environment variable is not configured,
// the testkit boots a shared local TiDB runtime and still creates an isolated
// database per test.
func StartIsolatedTiDB(t *testing.T) string {
	t.Helper()

	rootDSN := resolveTiDBRootDSN(t)
	return startTiDBDatabaseFn(t, rootDSN, "")
}

// StartIsolatedTiDBWithRootDSN returns a fresh empty TiDB database DSN on the
// provided root runtime and registers cleanup on the test.
func StartIsolatedTiDBWithRootDSN(t *testing.T, rootDSN string) string {
	t.Helper()
	return startTiDBDatabaseFn(t, rootDSN, "")
}

// StartIsolatedTiDBWithRootDSNAndCleanup returns a fresh empty TiDB database
// on the provided root runtime plus an explicit cleanup function. It is meant
// for package-level fixtures whose lifetime is longer than one test.
func StartIsolatedTiDBWithRootDSNAndCleanup(ctx context.Context, rootDSN string) (string, func(context.Context) error, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	return startTiDBDatabaseWithCleanupFn(ctx, rootDSN, "", tidbTemplatePrepareTimeout)
}

// StartClonedTiDBFromTemplate returns a fresh TiDB database cloned from the
// database named by templateDSN. Callers must keep the template alive and
// immutable until all clones have been created.
func StartClonedTiDBFromTemplate(t *testing.T, templateDSN string) string {
	t.Helper()

	cfg, err := mysql.ParseDSN(templateDSN)
	require.NoError(t, err)
	require.NotEmpty(t, cfg.DBName, "TiDB template DSN must name a database")

	templateDatabase := cfg.DBName
	cfg.DBName = ""
	configureTiDBTestDSN(cfg)
	return startTiDBDatabaseFn(t, cfg.FormatDSN(), templateDatabase)
}

// StartInitializedTiDB returns a fresh TiDB database DSN after applying the current schema.
func StartInitializedTiDB(t *testing.T) string {
	t.Helper()

	rootDSN := resolveTiDBRootDSN(t)
	return StartInitializedTiDBWithRootDSN(t, rootDSN)
}

// StartInitializedTiDBWithCleanup returns a fresh TiDB database DSN and an
// explicit cleanup function for package-level fixtures.
func StartInitializedTiDBWithCleanup(ctx context.Context) (string, func(context.Context) error, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	rootDSN, releaseRootDSN, err := resolveTiDBRootDSNWithCleanup(ctx)
	if err != nil {
		return "", nil, err
	}

	dsn, cleanupDatabase, err := startInitializedTiDBWithRootDSNAndCleanupFn(ctx, rootDSN)
	if err != nil {
		releaseRootDSN()
		return "", nil, err
	}

	cleanup := onceContextCleanup(func(cleanupCtx context.Context) error {
		if cleanupCtx == nil {
			cleanupCtx = context.Background()
		}
		err := cleanupDatabase(cleanupCtx)
		releaseRootDSN()
		return err
	})
	return dsn, cleanup, nil
}

func initializedTiDBTemplateDatabase(ctx context.Context, rootDSN string) (string, error) {
	configured := configuredTiDBTemplateDatabase()
	if configured != "" {
		tidbTemplateMu.Lock()
		if tidbPreparedTemplateRootDSN == rootDSN && tidbPreparedTemplateDatabase == configured {
			tidbTemplateMu.Unlock()
			return configured, nil
		}
		tidbTemplateMu.Unlock()

		ready, err := tidbTemplateReadyOnRootDSNFn(ctx, rootDSN, configured)
		if err != nil {
			return "", err
		}
		if ready {
			tidbTemplateMu.Lock()
			tidbPreparedTemplateRootDSN = rootDSN
			tidbPreparedTemplateDatabase = configured
			tidbTemplateMu.Unlock()
			return configured, nil
		}
	}

	tidbTemplateMu.Lock()
	defer tidbTemplateMu.Unlock()

	if tidbPreparedTemplateRootDSN == rootDSN && tidbPreparedTemplateDatabase != "" {
		return tidbPreparedTemplateDatabase, nil
	}

	templateDatabase, err := prepareMigratedTiDBTemplateFn(ctx, rootDSN)
	if err != nil {
		return "", err
	}

	tidbPreparedTemplateRootDSN = rootDSN
	tidbPreparedTemplateDatabase = templateDatabase
	return templateDatabase, nil
}

func tidbTemplateReadyOnRootDSN(ctx context.Context, rootDSN string, templateDatabase string) (bool, error) {
	adminDSN, err := tidbDSNWithoutDatabase(rootDSN)
	if err != nil {
		return false, err
	}

	adminDB, err := openSQLDBFn("mysql", adminDSN)
	if err != nil {
		return false, err
	}
	defer func() { _ = adminDB.Close() }()

	if err := adminDB.PingContext(ctx); err != nil {
		return false, err
	}

	exists, err := tidbDatabaseExistsFn(ctx, adminDB, templateDatabase)
	if err != nil || !exists {
		return false, err
	}

	current, err := tidbDatabaseSchemaCurrentFn(ctx, rootDSN, templateDatabase)
	if err != nil || !current {
		return false, err
	}

	sourceFingerprint, err := tidbTemplateSourceFingerprintFn()
	if err != nil {
		return false, err
	}
	return tidbTemplateFingerprintMatchesFn(ctx, adminDB, templateDatabase, sourceFingerprint)
}

func resetPreparedTiDBTemplateCache() {
	tidbTemplateMu.Lock()
	tidbPreparedTemplateRootDSN = ""
	tidbPreparedTemplateDatabase = ""
	tidbTemplateMu.Unlock()
	resetTiDBTemplateCloneMetadataCache()
}

// OpenTiDB opens and pings a TiDB database for tests.
func OpenTiDB(t *testing.T, dsn string) *sql.DB {
	t.Helper()

	ctx, cancel := context.WithTimeout(context.Background(), tidbTestDatabaseSetupTimeout)
	defer cancel()

	db, err := OpenTiDBContext(ctx, dsn)
	require.NoError(t, err)
	t.Cleanup(func() { _ = db.Close() })
	return db
}

// OpenTiDBContext opens and pings a TiDB database for package-level fixtures.
func OpenTiDBContext(ctx context.Context, dsn string) (*sql.DB, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}
	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	return db, nil
}

func startTiDBDatabase(t *testing.T, rootDSN string, templateDatabase string) string {
	t.Helper()

	ctx, cancel := context.WithTimeout(context.Background(), tidbTestDeadlineTimeout(t, tidbTemplatePrepareTimeout))
	defer cancel()

	operation := "create_database"
	if templateDatabase != "" {
		operation = "clone_database"
	}
	startedAt := time.Now()
	dsn, cleanup, err := startTiDBDatabaseWithCleanup(ctx, rootDSN, templateDatabase, tidbTestDeadlineTimeout(t, tidbTemplatePrepareTimeout))
	reportTiDBProvisionTiming(t, operation, time.Since(startedAt), err)
	require.NoError(t, err)

	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), tidbTestDatabaseSetupTimeout)
		defer cleanupCancel()
		require.NoError(t, cleanup(cleanupCtx))
	})
	return dsn
}

func reportTiDBProvisionTiming(t testing.TB, operation string, duration time.Duration, provisionErr error) {
	if os.Getenv(tidbTestTimingEnv) != "1" {
		return
	}
	status := "passed"
	if provisionErr != nil {
		status = "failed"
	}
	line := fmt.Sprintf(
		"[tidb:testkit:timing] test=%q operation=%s status=%s duration_ms=%d\n",
		t.Name(),
		operation,
		status,
		duration.Milliseconds(),
	)
	if timingFile := os.Getenv(tidbTestTimingFileEnv); timingFile != "" {
		file, err := os.OpenFile(timingFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
		if err == nil {
			_, _ = io.WriteString(file, line)
			_ = file.Close()
			return
		}
	}
	fmt.Fprint(os.Stderr, line)
}

func startTiDBDatabaseWithCleanup(ctx context.Context, rootDSN string, templateDatabase string, cloneSlotTimeout time.Duration) (string, func(context.Context) error, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return "", nil, err
	}

	adminDSN, err := tidbDSNWithoutDatabase(rootDSN)
	if err != nil {
		return "", nil, err
	}

	cfg, err := mysql.ParseDSN(rootDSN)
	if err != nil {
		return "", nil, err
	}
	configureTiDBTestDSN(cfg)

	databaseName := fmt.Sprintf("agentteam_tidb_test_%d_%d", time.Now().UnixNano(), atomic.AddUint64(&tidbTestDatabaseCounter, 1))

	cleanup := onceContextCleanup(func(cleanupCtx context.Context) error {
		if cleanupCtx == nil {
			cleanupCtx = context.Background()
		}
		return dropTiDBDatabaseWithAdminDSN(cleanupCtx, adminDSN, databaseName)
	})

	if templateDatabase == "" {
		err := retryTiDBProvisioning(ctx, databaseName, func(attemptCtx context.Context) error {
			createCtx, createCancel := context.WithTimeout(attemptCtx, tidbTestDatabaseSetupTimeout)
			defer createCancel()
			return withTiDBSetupAdminDB(createCtx, adminDSN, func(adminDB *sql.DB) error {
				return createTiDBDatabase(createCtx, adminDB, databaseName)
			})
		}, func(cleanupCtx context.Context) error {
			return dropTiDBDatabaseWithAdminDSN(cleanupCtx, adminDSN, databaseName)
		})
		if err != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), tidbTestDatabaseSetupTimeout)
			defer cleanupCancel()
			_ = cleanup(cleanupCtx)
			return "", nil, err
		}
	} else {
		if err := cloneTiDBDatabaseFromTemplateWithFreshAdmin(ctx, adminDSN, templateDatabase, databaseName, cloneSlotTimeout); err != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), tidbTestDatabaseSetupTimeout)
			defer cleanupCancel()
			_ = cleanup(cleanupCtx)
			return "", nil, err
		}
	}

	cfg.DBName = databaseName
	return cfg.FormatDSN(), cleanup, nil
}

func cloneTiDBDatabaseFromTemplateWithFreshAdmin(ctx context.Context, adminDSN string, templateDatabase string, databaseName string, cloneSlotTimeout time.Duration) error {
	return withSharedTiDBCloneSlotContextFn(ctx, cloneSlotTimeout, func() error {
		cacheNamespace := ""
		if templateDatabase == defaultTiDBTemplateDatabase && configuredTiDBTemplateDatabase() == "" {
			cacheNamespace = adminDSN
		}
		return retryTiDBProvisioning(ctx, databaseName, func(attemptCtx context.Context) error {
			cloneCtx, cloneCancel := context.WithTimeout(attemptCtx, tidbTestDatabaseSetupTimeout)
			defer cloneCancel()
			return withTiDBSetupAdminDB(cloneCtx, adminDSN, func(adminDB *sql.DB) error {
				return cloneTiDBDatabaseWithCacheFn(cloneCtx, adminDB, adminDSN, cacheNamespace, templateDatabase, databaseName)
			})
		}, func(cleanupCtx context.Context) error {
			return dropTiDBDatabaseWithAdminDSN(cleanupCtx, adminDSN, databaseName)
		})
	})
}

func retryTiDBProvisioning(ctx context.Context, databaseName string, attempt func(context.Context) error, cleanup func(context.Context) error) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if attempt == nil {
		return nil
	}

	var lastErr error
	for i := 1; i <= tidbProvisioningRetryAttempts; i++ {
		if err := ctx.Err(); err != nil {
			return err
		}

		err := attempt(ctx)
		if err == nil {
			return nil
		}
		lastErr = err
		if !isTransientTiDBProvisioningError(err) || i == tidbProvisioningRetryAttempts {
			return err
		}
		if cleanup != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), tidbTestDatabaseSetupTimeout)
			cleanupErr := cleanup(cleanupCtx)
			cleanupCancel()
			if cleanupErr != nil {
				return errors.Join(
					err,
					fmt.Errorf("cleanup partial tidb database %s before retry: %w", databaseName, cleanupErr),
				)
			}
		}
		if err := sleepContext(ctx, tidbProvisioningRetryDelay*time.Duration(i)); err != nil {
			return err
		}
	}
	return lastErr
}

func isTransientTiDBProvisioningError(err error) bool {
	if errors.Is(err, context.DeadlineExceeded) ||
		errors.Is(err, driver.ErrBadConn) ||
		errors.Is(err, mysql.ErrInvalidConn) ||
		errors.Is(err, io.EOF) {
		return true
	}
	var mysqlErr *mysql.MySQLError
	return errors.As(err, &mysqlErr) && mysqlErr.Number == tidbInformationSchemaOutOfDateErrorNumber
}

func dropTiDBDatabaseWithAdminDSN(ctx context.Context, adminDSN string, databaseName string) error {
	return retryTiDBProvisioning(ctx, databaseName, func(attemptCtx context.Context) error {
		dropCtx, dropCancel := context.WithTimeout(attemptCtx, tidbTestDatabaseSetupTimeout)
		defer dropCancel()
		return withTiDBSetupAdminDB(dropCtx, adminDSN, func(adminDB *sql.DB) error {
			return dropTiDBDatabaseFn(dropCtx, adminDB, databaseName)
		})
	}, nil)
}

func withTiDBSetupAdminDB(ctx context.Context, adminDSN string, fn func(*sql.DB) error) error {
	if fn == nil {
		return nil
	}
	adminDB, err := openSQLDBFn("mysql", adminDSN)
	if err != nil {
		return err
	}
	configureTiDBSetupSQLDB(adminDB)
	defer func() { _ = adminDB.Close() }()

	if err := adminDB.PingContext(ctx); err != nil {
		return err
	}
	return fn(adminDB)
}

func configureTiDBSetupSQLDB(db *sql.DB) {
	if db == nil {
		return
	}
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	db.SetConnMaxLifetime(tidbTestDatabaseSetupTimeout)
	db.SetConnMaxIdleTime(tidbTestDatabaseSetupTimeout / 2)
}

func onceContextCleanup(cleanup func(context.Context) error) func(context.Context) error {
	var once sync.Once
	var cleanupErr error
	return func(ctx context.Context) error {
		once.Do(func() {
			cleanupErr = cleanup(ctx)
		})
		return cleanupErr
	}
}

func tidbTestDeadlineTimeout(t *testing.T, fallback time.Duration) time.Duration {
	t.Helper()

	deadline, ok := t.Deadline()
	if !ok {
		return fallback
	}
	timeout := time.Until(deadline) - tidbTestDatabaseSetupTimeout
	if timeout > 0 {
		return timeout
	}
	return time.Nanosecond
}

func tidbDSNWithoutDatabase(dsn string) (string, error) {
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil {
		return "", fmt.Errorf("parse tidb dsn: %w", err)
	}
	configureTiDBTestDSN(cfg)
	cfg.DBName = ""
	return cfg.FormatDSN(), nil
}

func configureTiDBTestDSN(cfg *mysql.Config) {
	cfg.ParseTime = true
	if cfg.Loc == nil {
		cfg.Loc = time.UTC
	}
	if cfg.Params == nil {
		cfg.Params = make(map[string]string, 1)
	}
	if _, ok := cfg.Params["time_zone"]; !ok {
		cfg.Params["time_zone"] = "'+00:00'"
	}
}