package tidbruntime

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"io/fs"
	"math"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"testing/fstest"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/pressly/goose/v3"

	"github.com/tidbcloud/AgentTeam/internal/migrations"
)

const defaultTiDBTemplateDatabase = "agentteam_tidb_template_current"
const tidbTemplateFingerprintTable = "__agentteam_testkit_schema_fingerprint"
const tidbTemplateAdvisoryLockName = "agentteam_testkit_template_prepare"
const tidbTemplateAdvisoryLockReleaseTimeout = 5 * time.Second

type tidbTemplateCloneMetadata struct {
	tables      []tidbTemplateTableCloneMetadata
	foreignKeys []tidbTemplateForeignKeyCloneMetadata
}

type tidbTemplateTableCloneMetadata struct {
	name              string
	createStatement   string
	insertableColumns string
	rowCount          int64
}

type tidbTemplateForeignKeyCloneMetadata struct {
	tableName              string
	constraintName         string
	columns                []string
	referencedDatabaseName string
	referencedTableName    string
	referencedColumns      []string
	updateRule             string
	deleteRule             string
}

type tidbTemplateCloneMetadataCacheKey struct {
	cacheNamespace   string
	templateDatabase string
}

var (
	tidbTemplateCloneMetadataMu    sync.Mutex
	tidbTemplateCloneMetadataCache map[tidbTemplateCloneMetadataCacheKey]tidbTemplateCloneMetadata
)

var (
	openSQLDBFn                          = sql.Open
	tidbTemplateSourceFingerprintFn      = tidbTemplateSourceFingerprint
	tidbDatabaseSchemaCurrentFn          = tidbDatabaseSchemaCurrent
	tidbTemplateFingerprintMatchesFn     = tidbTemplateFingerprintMatches
	tidbDatabaseExistsFn                 = tidbDatabaseExists
	dropTiDBDatabaseFn                   = dropTiDBDatabase
	createTiDBDatabaseFn                 = createTiDBDatabase
	applyTiDBMigrationsToDatabaseFn      = applyTiDBMigrationsToDatabase
	writeTiDBTemplateFingerprintFn       = writeTiDBTemplateFingerprint
	inspectTiDBTemplateCloneMetadataFn   = inspectTiDBTemplateCloneMetadata
	withSharedTiDBAllCloneSlotsContextFn = withSharedTiDBAllCloneSlotsContext
)

func configuredTiDBTemplateDatabase() string {
	return strings.TrimSpace(os.Getenv("AGENTTEAM_TIDB_TEMPLATE_DATABASE"))
}

// PrepareMigratedTiDBTemplate ensures a reusable migrated TiDB template database exists
// for the provided root DSN and returns its database name.
func PrepareMigratedTiDBTemplate(ctx context.Context, rootDSN string) (string, error) {
	if strings.TrimSpace(rootDSN) == "" {
		return "", fmt.Errorf("tidb root dsn is required")
	}

	var templateDatabase string
	err := withSharedTiDBTemplateLockContext(ctx, tidbTemplatePrepareLockTimeout(ctx), func() error {
		var err error
		templateDatabase, err = prepareMigratedTiDBTemplateLocked(ctx, rootDSN)
		return err
	})
	if err != nil {
		return "", err
	}
	return templateDatabase, nil
}

func prepareMigratedTiDBTemplateLocked(ctx context.Context, rootDSN string) (string, error) {
	adminDSN, err := tidbDSNWithoutDatabase(rootDSN)
	if err != nil {
		return "", err
	}

	adminDB, err := openSQLDBFn("mysql", adminDSN)
	if err != nil {
		return "", fmt.Errorf("open tidb template admin db: %w", err)
	}
	defer func() { _ = adminDB.Close() }()

	if err := adminDB.PingContext(ctx); err != nil {
		return "", fmt.Errorf("ping tidb template admin db: %w", err)
	}

	var templateDatabase string
	err = withTiDBTemplateAdvisoryLock(ctx, adminDB, tidbTemplatePrepareLockTimeout(ctx), func() error {
		var prepareErr error
		templateDatabase, prepareErr = prepareMigratedTiDBTemplateOnAdminLocked(ctx, rootDSN, adminDB)
		return prepareErr
	})
	if err != nil {
		return "", err
	}
	return templateDatabase, nil
}

func prepareMigratedTiDBTemplateOnAdminLocked(ctx context.Context, rootDSN string, adminDB *sql.DB) (string, error) {
	templateDatabase := defaultTiDBTemplateDatabase
	sourceFingerprint, err := tidbTemplateSourceFingerprintFn()
	if err != nil {
		return "", err
	}
	current, err := tidbDatabaseSchemaCurrentFn(ctx, rootDSN, templateDatabase)
	rebuildTemplate := err != nil || !current
	if err == nil && current {
		matches, fingerprintErr := tidbTemplateFingerprintMatchesFn(ctx, adminDB, templateDatabase, sourceFingerprint)
		if fingerprintErr != nil {
			return "", fingerprintErr
		}
		if matches {
			return templateDatabase, nil
		}
		rebuildTemplate = true
	}
	if rebuildTemplate {
		err := withSharedTiDBAllCloneSlotsContextFn(ctx, tidbTemplatePrepareLockTimeout(ctx), func() error {
			cleanupTemplate := func(cleanupCtx context.Context) error {
				return dropTiDBDatabaseFn(cleanupCtx, adminDB, templateDatabase)
			}
			prepareErr := retryTiDBProvisioning(
				ctx,
				templateDatabase,
				func(attemptCtx context.Context) error {
					exists, existsErr := tidbDatabaseExistsFn(attemptCtx, adminDB, templateDatabase)
					if existsErr != nil {
						return existsErr
					}
					if exists {
						if dropErr := cleanupTemplate(attemptCtx); dropErr != nil {
							return dropErr
						}
					}

					if createErr := createTiDBDatabaseFn(attemptCtx, adminDB, templateDatabase); createErr != nil {
						return createErr
					}

					if migrationErr := applyTiDBMigrationsToDatabaseFn(
						attemptCtx,
						rootDSN,
						templateDatabase,
					); migrationErr != nil {
						return migrationErr
					}
					return writeTiDBTemplateFingerprintFn(
						attemptCtx,
						adminDB,
						templateDatabase,
						sourceFingerprint,
					)
				},
				cleanupTemplate,
			)
			if prepareErr != nil {
				cleanupCtx, cleanupCancel := context.WithTimeout(
					context.WithoutCancel(ctx),
					tidbTestDatabaseSetupTimeout,
				)
				defer cleanupCancel()
				if cleanupErr := cleanupTemplate(cleanupCtx); cleanupErr != nil {
					return errors.Join(prepareErr, fmt.Errorf(
						"cleanup failed tidb template %s: %w",
						templateDatabase,
						cleanupErr,
					))
				}
				return prepareErr
			}
			resetTiDBTemplateCloneMetadataCache()
			return nil
		})
		if err != nil {
			return "", err
		}
	}
	return templateDatabase, nil
}

func withTiDBTemplateAdvisoryLock(ctx context.Context, adminDB *sql.DB, timeout time.Duration, fn func() error) (err error) {
	if adminDB == nil {
		return fmt.Errorf("tidb template advisory lock requires an admin db")
	}
	if fn == nil {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}

	conn, err := adminDB.Conn(ctx)
	if err != nil {
		return fmt.Errorf("pin tidb template advisory lock session: %w", err)
	}
	defer func() {
		if closeErr := conn.Close(); closeErr != nil {
			err = errors.Join(err, fmt.Errorf("close tidb template advisory lock session: %w", closeErr))
		}
	}()

	waitSeconds := int(math.Ceil(timeout.Seconds()))
	if waitSeconds < 1 {
		waitSeconds = 1
	}
	var acquired sql.NullInt64
	if queryErr := conn.QueryRowContext(ctx, `SELECT GET_LOCK(?, ?)`, tidbTemplateAdvisoryLockName, waitSeconds).Scan(&acquired); queryErr != nil {
		return fmt.Errorf("acquire tidb template advisory lock: %w", queryErr)
	}
	if !acquired.Valid {
		return fmt.Errorf("acquire tidb template advisory lock: returned NULL")
	}
	if acquired.Int64 != 1 {
		if acquired.Int64 == 0 {
			return fmt.Errorf("acquire tidb template advisory lock: timed out")
		}
		return fmt.Errorf("acquire tidb template advisory lock: unexpected result %d", acquired.Int64)
	}

	defer func() {
		releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), tidbTemplateAdvisoryLockReleaseTimeout)
		defer cancel()

		var released sql.NullInt64
		if releaseErr := conn.QueryRowContext(releaseCtx, `SELECT RELEASE_LOCK(?)`, tidbTemplateAdvisoryLockName).Scan(&released); releaseErr != nil {
			err = errors.Join(err, fmt.Errorf("release tidb template advisory lock: %w", releaseErr))
			return
		}
		if !released.Valid {
			err = errors.Join(err, fmt.Errorf("release tidb template advisory lock: returned NULL"))
			return
		}
		if released.Int64 != 1 {
			err = errors.Join(err, fmt.Errorf("release tidb template advisory lock: current session did not release lock (result %d)", released.Int64))
		}
	}()

	return fn()
}

func tidbTemplatePrepareLockTimeout(ctx context.Context) time.Duration {
	if deadline, ok := ctx.Deadline(); ok {
		timeout := time.Until(deadline)
		if timeout > 0 {
			return timeout
		}
		return time.Nanosecond
	}
	return tidbTemplatePrepareTimeout
}

func applyTiDBMigrationsToDatabase(ctx context.Context, rootDSN string, databaseName string) error {
	dsn, err := tidbDSNWithDatabase(rootDSN, databaseName)
	if err != nil {
		return err
	}

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return fmt.Errorf("open tidb template database %s: %w", databaseName, err)
	}
	defer func() { _ = db.Close() }()

	if err := db.PingContext(ctx); err != nil {
		return fmt.Errorf("ping tidb template database %s: %w", databaseName, err)
	}

	if err := initializeTiDBSchema(ctx, db); err != nil {
		return fmt.Errorf("apply tidb template migrations: %w", err)
	}
	current, err := tidbSchemaCurrent(ctx, db)
	if err != nil {
		return fmt.Errorf("verify tidb template schema current: %w", err)
	}
	if !current {
		return fmt.Errorf("verify tidb template schema current: pending migrations remain")
	}
	return nil
}

func tidbTemplateSourceFingerprint() (string, error) {
	bootstrapSQL, err := tidbBootstrapSQL()
	if err != nil {
		return "", err
	}
	migrationFS, err := tidbMigrationFS(migrations.TiDBBootstrapBaselineVersion)
	if err != nil {
		return "", err
	}
	matches, err := fs.Glob(migrationFS, "*.sql")
	if err != nil {
		return "", fmt.Errorf("glob tidb migration fingerprint sources: %w", err)
	}
	sort.Strings(matches)

	hash := sha256.New()
	_, _ = hash.Write([]byte("bootstrap.sql\n"))
	_, _ = hash.Write([]byte(bootstrapSQL))
	for _, name := range matches {
		data, err := fs.ReadFile(migrationFS, name)
		if err != nil {
			return "", fmt.Errorf("read tidb migration fingerprint source %s: %w", name, err)
		}
		_, _ = hash.Write([]byte("\n" + name + "\n"))
		_, _ = hash.Write(data)
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}

func tidbTemplateFingerprintMatches(ctx context.Context, adminDB *sql.DB, templateDatabase string, want string) (bool, error) {
	var got string
	err := adminDB.QueryRowContext(ctx, `
SELECT fingerprint
FROM `+qualifiedTiDBTable(templateDatabase, tidbTemplateFingerprintTable)+`
WHERE id = 'source'
`).Scan(&got)
	if err == nil {
		return strings.TrimSpace(got) == strings.TrimSpace(want), nil
	}
	if errors.Is(err, sql.ErrNoRows) || isMissingTiDBTable(err) {
		return false, nil
	}
	return false, fmt.Errorf("read tidb template schema fingerprint: %w", err)
}

func writeTiDBTemplateFingerprint(ctx context.Context, adminDB *sql.DB, templateDatabase string, fingerprint string) error {
	if _, err := adminDB.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS `+qualifiedTiDBTable(templateDatabase, tidbTemplateFingerprintTable)+` (
  id varchar(64) NOT NULL,
  fingerprint varchar(128) NOT NULL,
  updated_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
  PRIMARY KEY (id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin
`); err != nil {
		return fmt.Errorf("create tidb template schema fingerprint table: %w", err)
	}
	if _, err := adminDB.ExecContext(ctx, `
INSERT INTO `+qualifiedTiDBTable(templateDatabase, tidbTemplateFingerprintTable)+` (id, fingerprint)
VALUES ('source', ?)
ON DUPLICATE KEY UPDATE fingerprint = VALUES(fingerprint), updated_at = CURRENT_TIMESTAMP
`, strings.TrimSpace(fingerprint)); err != nil {
		return fmt.Errorf("write tidb template schema fingerprint: %w", err)
	}
	return nil
}

func tidbDatabaseSchemaCurrent(ctx context.Context, rootDSN string, databaseName string) (bool, error) {
	dsn, err := tidbDSNWithDatabase(rootDSN, databaseName)
	if err != nil {
		return false, err
	}

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return false, err
	}
	defer func() { _ = db.Close() }()

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

func createTiDBDatabase(ctx context.Context, adminDB *sql.DB, databaseName string) error {
	if _, err := adminDB.ExecContext(ctx, "CREATE DATABASE "+quoteTiDBIdentifier(databaseName)); err != nil {
		return fmt.Errorf("create tidb database %s: %w", databaseName, err)
	}
	return nil
}

func dropTiDBDatabase(ctx context.Context, adminDB *sql.DB, databaseName string) error {
	if _, err := adminDB.ExecContext(ctx, "DROP DATABASE IF EXISTS "+quoteTiDBIdentifier(databaseName)); err != nil {
		return fmt.Errorf("drop tidb database %s: %w", databaseName, err)
	}
	return nil
}

func cloneTiDBDatabase(ctx context.Context, adminDB *sql.DB, templateDatabase string, databaseName string) error {
	return cloneTiDBDatabaseWithCache(ctx, adminDB, "", "", templateDatabase, databaseName)
}

func cloneTiDBDatabaseWithCache(ctx context.Context, adminDB *sql.DB, cloneAdminDSN string, cacheNamespace string, templateDatabase string, databaseName string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	metadata, err := loadTiDBTemplateCloneMetadata(ctx, adminDB, cacheNamespace, templateDatabase)
	if err != nil {
		return err
	}
	if err := createTiDBDatabase(ctx, adminDB, databaseName); err != nil {
		return err
	}

	sessionDB := adminDB
	if strings.TrimSpace(cloneAdminDSN) != "" {
		dedicatedDB, err := sql.Open("mysql", cloneAdminDSN)
		if err != nil {
			return fmt.Errorf("open tidb clone session db: %w", err)
		}
		configureTiDBSetupSQLDB(dedicatedDB)
		defer func() { _ = dedicatedDB.Close() }()
		sessionDB = dedicatedDB
	}

	conn, err := sessionDB.Conn(ctx)
	if err != nil {
		return fmt.Errorf("pin tidb clone session: %w", err)
	}
	defer func() { _ = conn.Close() }()

	foreignKeysByTable := tidbTemplateForeignKeysByTable(metadata.foreignKeys)
	return withTiDBForeignKeyChecksDisabled(ctx, conn, func() error {
		for _, table := range metadata.tables {
			if err := cloneTiDBTableStructure(ctx, conn, templateDatabase, databaseName, table, foreignKeysByTable[table.name]); err != nil {
				return err
			}
		}
		for _, table := range metadata.tables {
			if table.rowCount == 0 {
				continue
			}
			if err := cloneTiDBTableRowsWithColumns(ctx, conn, templateDatabase, databaseName, table.name, table.insertableColumns); err != nil {
				return err
			}
		}
		return nil
	})
}

func withTiDBForeignKeyChecksDisabled(ctx context.Context, conn *sql.Conn, fn func() error) (err error) {
	if conn == nil {
		return fmt.Errorf("tidb clone session is nil")
	}
	var previousForeignKeyChecks int
	if err := conn.QueryRowContext(ctx, "SELECT @@FOREIGN_KEY_CHECKS").Scan(&previousForeignKeyChecks); err != nil {
		return fmt.Errorf("read tidb foreign key checks: %w", err)
	}
	if previousForeignKeyChecks != 0 && previousForeignKeyChecks != 1 {
		return fmt.Errorf("read tidb foreign key checks: unexpected value %d", previousForeignKeyChecks)
	}
	if _, err := conn.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS=0"); err != nil {
		return fmt.Errorf("disable tidb foreign key checks: %w", err)
	}
	defer func() {
		restoreCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
		defer cancel()
		restoreQuery := "SET FOREIGN_KEY_CHECKS=1"
		if previousForeignKeyChecks == 0 {
			restoreQuery = "SET FOREIGN_KEY_CHECKS=0"
		}
		if _, restoreErr := conn.ExecContext(restoreCtx, restoreQuery); restoreErr != nil {
			err = errors.Join(err, fmt.Errorf("restore tidb foreign key checks: %w", restoreErr))
		}
	}()

	return fn()
}

func tidbTemplateForeignKeysByTable(foreignKeys []tidbTemplateForeignKeyCloneMetadata) map[string][]tidbTemplateForeignKeyCloneMetadata {
	byTable := make(map[string][]tidbTemplateForeignKeyCloneMetadata)
	for _, foreignKey := range foreignKeys {
		byTable[foreignKey.tableName] = append(byTable[foreignKey.tableName], foreignKey)
	}
	return byTable
}

func loadTiDBTemplateCloneMetadata(ctx context.Context, adminDB *sql.DB, cacheNamespace string, templateDatabase string) (tidbTemplateCloneMetadata, error) {
	key := tidbTemplateCloneMetadataCacheKey{
		cacheNamespace:   strings.TrimSpace(cacheNamespace),
		templateDatabase: templateDatabase,
	}
	if key.cacheNamespace != "" {
		tidbTemplateCloneMetadataMu.Lock()
		if cached, ok := tidbTemplateCloneMetadataCache[key]; ok {
			tidbTemplateCloneMetadataMu.Unlock()
			return cached, nil
		}
		tidbTemplateCloneMetadataMu.Unlock()
	}

	metadata, err := inspectTiDBTemplateCloneMetadataFn(ctx, adminDB, templateDatabase)
	if err != nil {
		return tidbTemplateCloneMetadata{}, err
	}
	if key.cacheNamespace == "" {
		return metadata, nil
	}

	tidbTemplateCloneMetadataMu.Lock()
	defer tidbTemplateCloneMetadataMu.Unlock()

	if tidbTemplateCloneMetadataCache == nil {
		tidbTemplateCloneMetadataCache = make(map[tidbTemplateCloneMetadataCacheKey]tidbTemplateCloneMetadata)
	}
	if cached, ok := tidbTemplateCloneMetadataCache[key]; ok {
		return cached, nil
	}
	tidbTemplateCloneMetadataCache[key] = metadata
	return metadata, nil
}

func inspectTiDBTemplateCloneMetadata(ctx context.Context, adminDB *sql.DB, templateDatabase string) (tidbTemplateCloneMetadata, error) {
	exists, err := tidbDatabaseExists(ctx, adminDB, templateDatabase)
	if err != nil {
		return tidbTemplateCloneMetadata{}, err
	}
	if !exists {
		return tidbTemplateCloneMetadata{}, fmt.Errorf("tidb template database %s does not exist", templateDatabase)
	}

	tables, err := listTiDBTemplateTables(ctx, adminDB, templateDatabase)
	if err != nil {
		return tidbTemplateCloneMetadata{}, err
	}

	metadata := tidbTemplateCloneMetadata{
		tables: make([]tidbTemplateTableCloneMetadata, 0, len(tables)),
	}
	for _, table := range tables {
		createStatement, err := showTiDBCreateTableStatement(ctx, adminDB, templateDatabase, table)
		if err != nil {
			return tidbTemplateCloneMetadata{}, err
		}

		columns, err := listTiDBInsertableColumns(ctx, adminDB, templateDatabase, table)
		if err != nil {
			return tidbTemplateCloneMetadata{}, err
		}
		if len(columns) == 0 {
			return tidbTemplateCloneMetadata{}, fmt.Errorf("list tidb insertable columns %s.%s: no insertable columns found", templateDatabase, table)
		}

		rowCount, err := countTiDBTableRows(ctx, adminDB, templateDatabase, table)
		if err != nil {
			return tidbTemplateCloneMetadata{}, err
		}
		metadata.tables = append(metadata.tables, tidbTemplateTableCloneMetadata{
			name:              table,
			createStatement:   createStatement,
			insertableColumns: joinQuotedTiDBIdentifiers(columns),
			rowCount:          rowCount,
		})
	}
	foreignKeys, err := listTiDBForeignKeys(ctx, adminDB, templateDatabase)
	if err != nil {
		return tidbTemplateCloneMetadata{}, err
	}
	metadata.foreignKeys = foreignKeys
	return metadata, nil
}

func resetTiDBTemplateCloneMetadataCache() {
	tidbTemplateCloneMetadataMu.Lock()
	tidbTemplateCloneMetadataCache = nil
	tidbTemplateCloneMetadataMu.Unlock()
}

func listTiDBTemplateTables(ctx context.Context, adminDB *sql.DB, databaseName string) ([]string, error) {
	rows, err := adminDB.QueryContext(ctx, `
SELECT table_name
FROM information_schema.tables
WHERE table_schema = ?
  AND table_type = 'BASE TABLE'
ORDER BY table_name ASC
`, databaseName)
	if err != nil {
		return nil, fmt.Errorf("list tidb template tables: %w", err)
	}
	defer rows.Close()

	tables := make([]string, 0)
	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err != nil {
			return nil, fmt.Errorf("scan tidb template table: %w", err)
		}
		tables = append(tables, tableName)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate tidb template tables: %w", err)
	}
	return tables, nil
}

type tidbSQLExecutor interface {
	ExecContext(context.Context, string, ...any) (sql.Result, error)
	QueryContext(context.Context, string, ...any) (*sql.Rows, error)
	QueryRowContext(context.Context, string, ...any) *sql.Row
}

func showTiDBCreateTableStatement(ctx context.Context, adminDB tidbSQLExecutor, databaseName string, tableName string) (string, error) {
	var table string
	var createStatement string
	if err := adminDB.QueryRowContext(ctx, "SHOW CREATE TABLE "+qualifiedTiDBTable(databaseName, tableName)).Scan(&table, &createStatement); err != nil {
		return "", fmt.Errorf("show tidb create table %s.%s: %w", databaseName, tableName, err)
	}
	return createStatement, nil
}

func cloneTiDBTableStructure(ctx context.Context, adminDB tidbSQLExecutor, templateDatabase string, databaseName string, table tidbTemplateTableCloneMetadata, foreignKeys []tidbTemplateForeignKeyCloneMetadata) error {
	query, err := buildTiDBCreateTableStatement(templateDatabase, databaseName, table, foreignKeys)
	if err != nil {
		return err
	}
	if _, err := adminDB.ExecContext(ctx, query); err != nil {
		return fmt.Errorf("clone tidb table structure %s.%s: %w", templateDatabase, table.name, err)
	}
	return nil
}

func buildTiDBCreateTableStatement(templateDatabase string, databaseName string, table tidbTemplateTableCloneMetadata, foreignKeys []tidbTemplateForeignKeyCloneMetadata) (string, error) {
	source := strings.TrimSpace(table.createStatement)
	if source == "" {
		return "", fmt.Errorf("build tidb create table statement %s: empty create statement", table.name)
	}

	unqualifiedPrefix := "CREATE TABLE " + quoteTiDBIdentifier(table.name)
	qualifiedPrefix := "CREATE TABLE " + qualifiedTiDBTable(templateDatabase, table.name)

	var rest string
	switch {
	case strings.HasPrefix(source, unqualifiedPrefix):
		rest = strings.TrimPrefix(source, unqualifiedPrefix)
	case strings.HasPrefix(source, qualifiedPrefix):
		rest = strings.TrimPrefix(source, qualifiedPrefix)
	default:
		return "", fmt.Errorf("build tidb create table statement %s: unexpected create statement", table.name)
	}

	statement := "CREATE TABLE " + qualifiedTiDBTable(databaseName, table.name) + rest
	return rewriteTiDBCreateTableReferences(statement, templateDatabase, databaseName, foreignKeys), nil
}

func rewriteTiDBCreateTableReferences(statement string, templateDatabase string, databaseName string, foreignKeys []tidbTemplateForeignKeyCloneMetadata) string {
	for _, foreignKey := range foreignKeys {
		referencedDatabaseName := foreignKey.referencedDatabaseName
		if referencedDatabaseName == "" || referencedDatabaseName == templateDatabase {
			referencedDatabaseName = databaseName
		}
		targetReference := "REFERENCES " + qualifiedTiDBTable(referencedDatabaseName, foreignKey.referencedTableName)
		for _, sourceReference := range []string{
			"REFERENCES " + quoteTiDBIdentifier(foreignKey.referencedTableName),
			"REFERENCES " + qualifiedTiDBTable(templateDatabase, foreignKey.referencedTableName),
		} {
			if sourceReference != "REFERENCES " && sourceReference != targetReference {
				statement = strings.ReplaceAll(statement, sourceReference, targetReference)
			}
		}
		if foreignKey.referencedDatabaseName != "" && foreignKey.referencedDatabaseName != templateDatabase {
			sourceReference := "REFERENCES " + qualifiedTiDBTable(foreignKey.referencedDatabaseName, foreignKey.referencedTableName)
			if sourceReference != targetReference {
				statement = strings.ReplaceAll(statement, sourceReference, targetReference)
			}
		}
	}
	return statement
}

func cloneTiDBTableRowsWithColumns(ctx context.Context, adminDB tidbSQLExecutor, templateDatabase string, databaseName string, tableName string, columnList string) error {
	query := "INSERT INTO " + qualifiedTiDBTable(databaseName, tableName) +
		" (" + columnList + ")" +
		" SELECT " + columnList +
		" FROM " + qualifiedTiDBTable(templateDatabase, tableName)
	if _, err := adminDB.ExecContext(ctx, query); err != nil {
		return fmt.Errorf("clone tidb table rows %s.%s: %w", templateDatabase, tableName, err)
	}
	return nil
}

func listTiDBForeignKeys(ctx context.Context, adminDB *sql.DB, databaseName string) ([]tidbTemplateForeignKeyCloneMetadata, error) {
	rows, err := adminDB.QueryContext(ctx, `
SELECT
  rc.table_name,
  rc.constraint_name,
  kcu.column_name,
  kcu.referenced_table_schema,
  rc.referenced_table_name,
  kcu.referenced_column_name,
  rc.update_rule,
  rc.delete_rule
FROM information_schema.referential_constraints rc
JOIN information_schema.key_column_usage kcu
  ON kcu.constraint_schema = rc.constraint_schema
 AND kcu.constraint_name = rc.constraint_name
 AND kcu.table_schema = rc.constraint_schema
 AND kcu.table_name = rc.table_name
WHERE rc.constraint_schema = ?
  AND kcu.referenced_table_name IS NOT NULL
ORDER BY rc.table_name ASC, rc.constraint_name ASC, kcu.ordinal_position ASC
`, databaseName)
	if err != nil {
		return nil, fmt.Errorf("list tidb foreign keys %s: %w", databaseName, err)
	}
	defer rows.Close()

	foreignKeys := make([]tidbTemplateForeignKeyCloneMetadata, 0)
	for rows.Next() {
		var tableName string
		var constraintName string
		var columnName string
		var referencedDatabaseName string
		var referencedTableName string
		var referencedColumnName string
		var updateRule string
		var deleteRule string
		if err := rows.Scan(
			&tableName,
			&constraintName,
			&columnName,
			&referencedDatabaseName,
			&referencedTableName,
			&referencedColumnName,
			&updateRule,
			&deleteRule,
		); err != nil {
			return nil, fmt.Errorf("scan tidb foreign key %s: %w", databaseName, err)
		}

		if len(foreignKeys) == 0 ||
			foreignKeys[len(foreignKeys)-1].tableName != tableName ||
			foreignKeys[len(foreignKeys)-1].constraintName != constraintName {
			foreignKeys = append(foreignKeys, tidbTemplateForeignKeyCloneMetadata{
				tableName:              tableName,
				constraintName:         constraintName,
				referencedDatabaseName: referencedDatabaseName,
				referencedTableName:    referencedTableName,
				updateRule:             updateRule,
				deleteRule:             deleteRule,
			})
		}
		index := len(foreignKeys) - 1
		foreignKeys[index].columns = append(foreignKeys[index].columns, columnName)
		foreignKeys[index].referencedColumns = append(foreignKeys[index].referencedColumns, referencedColumnName)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate tidb foreign keys %s: %w", databaseName, err)
	}
	return foreignKeys, nil
}

func listTiDBInsertableColumns(ctx context.Context, adminDB tidbSQLExecutor, databaseName string, tableName string) ([]string, error) {
	rows, err := adminDB.QueryContext(ctx, `
SELECT column_name
FROM information_schema.columns
WHERE table_schema = ?
  AND table_name = ?
  AND COALESCE(generation_expression, '') = ''
ORDER BY ordinal_position ASC
`, databaseName, tableName)
	if err != nil {
		return nil, fmt.Errorf("list tidb insertable columns %s.%s: %w", databaseName, tableName, err)
	}
	defer rows.Close()

	columns := make([]string, 0)
	for rows.Next() {
		var columnName string
		if err := rows.Scan(&columnName); err != nil {
			return nil, fmt.Errorf("scan tidb insertable column %s.%s: %w", databaseName, tableName, err)
		}
		columns = append(columns, columnName)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate tidb insertable columns %s.%s: %w", databaseName, tableName, err)
	}
	return columns, nil
}

func countTiDBTableRows(ctx context.Context, adminDB tidbSQLExecutor, databaseName string, tableName string) (int64, error) {
	row := adminDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+qualifiedTiDBTable(databaseName, tableName))
	var count int64
	if err := row.Scan(&count); err != nil {
		return 0, fmt.Errorf("count tidb table rows %s.%s: %w", databaseName, tableName, err)
	}
	if count < 0 {
		return 0, fmt.Errorf("count tidb table rows %s.%s: invalid row count", databaseName, tableName)
	}
	return count, nil
}

func tidbDatabaseExists(ctx context.Context, adminDB *sql.DB, databaseName string) (bool, error) {
	var count int
	if err := adminDB.QueryRowContext(ctx, `
SELECT COUNT(*)
FROM information_schema.schemata
WHERE schema_name = ?
`, databaseName).Scan(&count); err != nil {
		return false, fmt.Errorf("check tidb database %s exists: %w", databaseName, err)
	}
	return count > 0, nil
}

func initializeTiDBSchema(ctx context.Context, db *sql.DB) error {
	if db == nil {
		return fmt.Errorf("tidb db is nil")
	}

	bootstrapSQL, err := tidbBootstrapSQL()
	if err != nil {
		return err
	}
	if err := migrations.ExecSQLStatements(ctx, db, bootstrapSQL); err != nil {
		return fmt.Errorf("exec tidb bootstrap schema: %w", err)
	}
	if err := migrations.SeedTiDBGooseBaseline(ctx, db, migrations.TiDBBootstrapBaselineVersion); err != nil {
		return err
	}

	migrationFS, err := tidbMigrationFS(migrations.TiDBBootstrapBaselineVersion)
	if err != nil {
		return err
	}

	migrations.GooseMu.Lock()
	defer migrations.GooseMu.Unlock()

	goose.SetBaseFS(migrationFS)
	if err := goose.SetDialect(string(goose.DialectTiDB)); err != nil {
		return fmt.Errorf("set goose dialect: %w", err)
	}
	if _, err := goose.EnsureDBVersionContext(ctx, db); err != nil {
		return fmt.Errorf("ensure goose version table: %w", err)
	}
	if err := goose.UpContext(ctx, db, "."); err != nil {
		return fmt.Errorf("run goose up: %w", err)
	}
	return nil
}

func tidbSchemaCurrent(ctx context.Context, db *sql.DB) (bool, error) {
	state, err := inspectTiDBMigrationState(ctx, db)
	if err != nil {
		return false, err
	}
	return !state.Pending, nil
}

type tidbMigrationState struct {
	CurrentVersion int64
	TargetVersion  int64
	Pending        bool
}

func inspectTiDBMigrationState(ctx context.Context, db *sql.DB) (tidbMigrationState, error) {
	if db == nil {
		return tidbMigrationState{}, fmt.Errorf("tidb db is nil")
	}

	provider, err := newTiDBMigrationProvider(db)
	if err != nil {
		return tidbMigrationState{}, err
	}

	currentVersion, targetVersion, err := provider.GetVersions(ctx)
	if err != nil {
		if isMissingTiDBGooseVersionTable(err) {
			targetVersion = latestTiDBSourceVersion(provider)
			return tidbMigrationState{
				CurrentVersion: 0,
				TargetVersion:  targetVersion,
				Pending:        targetVersion > 0,
			}, nil
		}
		return tidbMigrationState{}, fmt.Errorf("read tidb migration versions: %w", err)
	}
	pending, err := provider.HasPending(ctx)
	if err != nil {
		return tidbMigrationState{}, fmt.Errorf("check tidb pending migrations: %w", err)
	}
	return tidbMigrationState{
		CurrentVersion: currentVersion,
		TargetVersion:  targetVersion,
		Pending:        pending,
	}, nil
}

func newTiDBMigrationProvider(db *sql.DB) (*goose.Provider, error) {
	migrationFS, err := tidbMigrationFS(migrations.TiDBBootstrapBaselineVersion)
	if err != nil {
		return nil, err
	}

	provider, err := goose.NewProvider(goose.DialectTiDB, db, migrationFS)
	if err != nil {
		return nil, fmt.Errorf("create tidb migration provider: %w", err)
	}
	return provider, nil
}

func tidbMigrationFS(minVersion int64) (fs.FS, error) {
	dir, err := tidbMigrationsDir()
	if err != nil {
		return nil, err
	}

	merged := fstest.MapFS{}
	if err := mergeTiDBMigrationFiles(merged, os.DirFS(dir), "*.sql", minVersion); err != nil {
		return nil, err
	}
	return merged, nil
}

func mergeTiDBMigrationFiles(dst fstest.MapFS, source fs.FS, pattern string, minVersion int64) error {
	matches, err := fs.Glob(source, pattern)
	if err != nil {
		return fmt.Errorf("glob tidb migrations %q: %w", pattern, err)
	}
	for _, name := range matches {
		version, err := goose.NumericComponent(path.Base(name))
		if err != nil {
			return fmt.Errorf("parse tidb migration version %s: %w", name, err)
		}
		if version < minVersion {
			continue
		}

		data, err := fs.ReadFile(source, name)
		if err != nil {
			return fmt.Errorf("read tidb migration %s: %w", name, err)
		}
		base := path.Base(name)
		if _, exists := dst[base]; exists {
			return fmt.Errorf("duplicate tidb migration file %s", base)
		}
		dst[base] = &fstest.MapFile{Data: data, Mode: 0o644}
	}
	return nil
}

func tidbBootstrapSQL() (string, error) {
	path, err := tidbBootstrapPath()
	if err != nil {
		return "", err
	}

	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("read tidb bootstrap sql: %w", err)
	}
	return string(data), nil
}

func tidbBootstrapPath() (string, error) {
	_, currentFile, _, ok := runtime.Caller(0)
	if !ok {
		return "", fmt.Errorf("resolve tidb bootstrap path: runtime caller unavailable")
	}

	path := filepath.Join(filepath.Dir(currentFile), "..", "..", "..", "tidb", "bootstrap.sql")
	if _, err := os.Stat(path); err != nil {
		return "", fmt.Errorf("resolve tidb bootstrap path: %w", err)
	}
	return path, nil
}

func tidbMigrationsDir() (string, error) {
	_, currentFile, _, ok := runtime.Caller(0)
	if !ok {
		return "", fmt.Errorf("resolve tidb migrations dir: runtime caller unavailable")
	}

	dir := filepath.Join(filepath.Dir(currentFile), "..", "..", "..", "tidb", "migrations")
	if _, err := os.Stat(dir); err != nil {
		return "", fmt.Errorf("resolve tidb migrations dir: %w", err)
	}
	return dir, nil
}

func latestTiDBSourceVersion(provider *goose.Provider) int64 {
	sources := provider.ListSources()
	if len(sources) == 0 {
		return 0
	}
	return sources[len(sources)-1].Version
}

func isMissingTiDBGooseVersionTable(err error) bool {
	if isMissingTiDBTable(err) {
		return true
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "goose_db_version") && strings.Contains(message, "doesn't exist")
}

func isMissingTiDBTable(err error) bool {
	var mysqlErr *mysql.MySQLError
	if errors.As(err, &mysqlErr) && mysqlErr.Number == 1146 {
		return true
	}
	message := strings.ToLower(err.Error())
	return strings.Contains(message, "doesn't exist")
}

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

func qualifiedTiDBTable(databaseName string, tableName string) string {
	return quoteTiDBIdentifier(databaseName) + "." + quoteTiDBIdentifier(tableName)
}

func quoteTiDBIdentifier(identifier string) string {
	return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
}

func joinQuotedTiDBIdentifiers(identifiers []string) string {
	if len(identifiers) == 0 {
		return ""
	}

	var builder strings.Builder
	for index, identifier := range identifiers {
		if index > 0 {
			builder.WriteString(", ")
		}
		builder.WriteString(quoteTiDBIdentifier(identifier))
	}
	return builder.String()
}