package testkit

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

	"github.com/stretchr/testify/require"

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

// TiDBCoreDBFixture owns one shared TiDB runtime for a test package while
// provisioning an isolated migrated database for each test DB handle.
type TiDBCoreDBFixture struct {
	runtime                    TiDBTestRuntime
	runtimeFixture             *TiDBRuntimeFixture
	once                       sync.Once
	serializedRollbackMu       sync.Mutex
	serializedRollbackOnce     sync.Once
	serializedRollbackSQLDB    *sql.DB
	serializedRollbackCleanup  func(context.Context) error
	serializedRollbackErr      error
	serializedCommittedMu      sync.Mutex
	serializedCommittedOnce    sync.Once
	serializedCommittedSQLDB   *sql.DB
	serializedCommittedCleanup func(context.Context) error
	serializedCommittedErr     error
}

type sharedTiDBRuntimeRootState struct {
	fixture *TiDBRuntimeFixture
	users   int
	closing bool
	done    chan struct{}
}

var (
	sharedTiDBRuntimeRootMu sync.Mutex
	sharedTiDBRuntimeRoots  = make(map[TiDBTestRuntime]*sharedTiDBRuntimeRootState)

	packageTiDBRuntimeRootMu sync.Mutex
	packageTiDBRuntimeRoots  = make(map[TiDBTestRuntime]*TiDBRuntimeFixture)
)

func NewTiDBCoreDBFixture(runtime TiDBTestRuntime) *TiDBCoreDBFixture {
	return &TiDBCoreDBFixture{runtime: runtime}
}

func (fixture *TiDBCoreDBFixture) runtimeOrDefault() TiDBTestRuntime {
	if fixture != nil && fixture.runtime != "" {
		return fixture.runtime
	}
	return TiDBTestRuntimeUnistore
}

func (fixture *TiDBCoreDBFixture) sharedRuntimeFixture() *TiDBRuntimeFixture {
	if fixture == nil {
		return nil
	}
	fixture.once.Do(func() {
		fixture.runtimeFixture = NewTiDBRuntimeFixture(fixture.runtimeOrDefault())
	})
	return fixture.runtimeFixture
}

// RootDSN returns the package fixture's root TiDB DSN for tests that need to
// open a migrated database without wrapping it in the rollback fixture.
func (fixture *TiDBCoreDBFixture) RootDSN(t testing.TB) string {
	t.Helper()
	if fixture == nil {
		t.Fatal("tidb coredb fixture is nil")
	}
	return fixture.sharedRuntimeFixture().RootDSN(t)
}

// RollbackDB returns a DB scoped to a test transaction in the package fixture.
func (fixture *TiDBCoreDBFixture) RollbackDB(t *testing.T, ctx context.Context) coredb.DB {
	t.Helper()
	if fixture == nil {
		t.Fatal("tidb coredb fixture is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	rootDSN := fixture.sharedRuntimeFixture().RootDSN(t)
	dsn := StartInitializedTiDBWithRootDSN(t, rootDSN)
	openCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
	defer cancel()
	sqlDB, err := OpenTiDBContext(openCtx, dsn)
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(2)
	sqlDB.SetMaxIdleConns(2)
	t.Cleanup(func() {
		if err := sqlDB.Close(); err != nil {
			t.Errorf("close tidb rollback db: %v", err)
		}
	})
	return BeginTiDBRollbackDB(t, ctx, coredb.WrapSQLDB(sqlDB))
}

// SerializedRollbackDB returns a transaction on one package-level migrated
// database. Calls are serialized for the full test lifetime, and every test
// transaction rolls back before the next caller can use the database.
//
// This fixture is only safe for transaction-scoped state. Callers must not use
// t.Parallel, DDL, TiDB global variables, failpoints, process environment, or
// runtime lifecycle operations: rollback cannot isolate those effects, and a
// caller that pauses via t.Parallel after acquiring the lock can deadlock the
// package. Use RollbackDB or a dedicated runtime when those operations matter.
func (fixture *TiDBCoreDBFixture) SerializedRollbackDB(t *testing.T, ctx context.Context) coredb.DB {
	t.Helper()
	if fixture == nil {
		t.Fatal("tidb coredb fixture is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	fixture.serializedRollbackMu.Lock()
	// BeginTiDBRollbackDB registers its rollback after this cleanup, so testing's
	// LIFO cleanup order rolls the transaction back before releasing the lock.
	t.Cleanup(fixture.serializedRollbackMu.Unlock)
	fixture.serializedRollbackOnce.Do(func() {
		rootDSN := fixture.sharedRuntimeFixture().RootDSN(t)
		setupCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
		defer cancel()

		dsn, cleanup, err := StartInitializedTiDBWithRootDSNAndCleanup(setupCtx, rootDSN)
		if err != nil {
			fixture.serializedRollbackErr = err
			return
		}
		sqlDB, err := OpenTiDBContext(setupCtx, dsn)
		if err != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
			defer cleanupCancel()
			fixture.serializedRollbackErr = errors.Join(err, cleanup(cleanupCtx))
			return
		}
		fixture.serializedRollbackCleanup = cleanup
		fixture.serializedRollbackSQLDB = sqlDB
		fixture.serializedRollbackSQLDB.SetMaxOpenConns(2)
		fixture.serializedRollbackSQLDB.SetMaxIdleConns(2)
	})

	require.NoError(t, fixture.serializedRollbackErr)
	require.NotNil(t, fixture.serializedRollbackSQLDB)
	return BeginTiDBRollbackDB(t, ctx, coredb.WrapSQLDB(fixture.serializedRollbackSQLDB))
}

// SerializedCommittedDB returns one package-level migrated database with a
// real connection pool. Calls are serialized for the full test lifetime, and
// reset must remove every row owned by the previous test while preserving
// migrated seed data. This is for tests whose multi-connection or commit
// semantics cannot run inside SerializedRollbackDB.
//
// The reset callback is the owning package's isolation boundary: keep it
// bounded to explicit test-owned roots, and fail if those roots do not cascade
// cleanly. Callers must not use t.Parallel, DDL, TiDB global variables,
// failpoints, process environment, or runtime lifecycle operations.
func (fixture *TiDBCoreDBFixture) SerializedCommittedDB(
	t *testing.T,
	ctx context.Context,
	reset func(context.Context, coredb.DB) error,
) coredb.DB {
	t.Helper()
	if fixture == nil {
		t.Fatal("tidb coredb fixture is nil")
	}
	if reset == nil {
		t.Fatal("serialized committed database reset is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	fixture.serializedCommittedMu.Lock()
	t.Cleanup(fixture.serializedCommittedMu.Unlock)
	fixture.serializedCommittedOnce.Do(func() {
		rootDSN := fixture.sharedRuntimeFixture().RootDSN(t)
		setupCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
		defer cancel()

		dsn, cleanup, err := StartInitializedTiDBWithRootDSNAndCleanup(setupCtx, rootDSN)
		if err != nil {
			fixture.serializedCommittedErr = err
			return
		}
		sqlDB, err := OpenTiDBContext(setupCtx, dsn)
		if err != nil {
			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 2*time.Minute)
			defer cleanupCancel()
			fixture.serializedCommittedErr = errors.Join(err, cleanup(cleanupCtx))
			return
		}
		fixture.serializedCommittedCleanup = cleanup
		fixture.serializedCommittedSQLDB = sqlDB
		fixture.serializedCommittedSQLDB.SetMaxOpenConns(8)
		fixture.serializedCommittedSQLDB.SetMaxIdleConns(8)
	})

	require.NoError(t, fixture.serializedCommittedErr)
	require.NotNil(t, fixture.serializedCommittedSQLDB)
	db := coredb.WrapSQLDB(fixture.serializedCommittedSQLDB)
	require.NoError(t, reset(ctx, db))
	return db
}

// Cleanup releases the package-level TiDB fixture.
func (fixture *TiDBCoreDBFixture) Cleanup(ctx context.Context) error {
	if fixture == nil || fixture.sharedRuntimeFixture() == nil {
		return nil
	}

	var cleanupErr error
	if fixture.serializedRollbackSQLDB != nil {
		cleanupErr = errors.Join(cleanupErr, fixture.serializedRollbackSQLDB.Close())
	}
	if fixture.serializedRollbackCleanup != nil {
		cleanupErr = errors.Join(cleanupErr, fixture.serializedRollbackCleanup(ctx))
	}
	if fixture.serializedCommittedSQLDB != nil {
		cleanupErr = errors.Join(cleanupErr, fixture.serializedCommittedSQLDB.Close())
	}
	if fixture.serializedCommittedCleanup != nil {
		cleanupErr = errors.Join(cleanupErr, fixture.serializedCommittedCleanup(ctx))
	}
	return errors.Join(cleanupErr, fixture.runtimeFixture.Cleanup(ctx))
}

// RunTestsWithTiDBCoreDBFixture runs a test package and cleans up its TiDB fixture.
func RunTestsWithTiDBCoreDBFixture(m *testing.M, name string, fixture *TiDBCoreDBFixture) int {
	code := m.Run()
	if fixture == nil {
		return code
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	if err := fixture.Cleanup(ctx); err != nil {
		fmt.Fprintf(os.Stderr, "cleanup %s tidb fixture: %v\n", name, err)
		if code == 0 {
			code = 1
		}
	}
	return code
}

// StartInitializedTiDBWithSharedRuntime returns a migrated isolated database
// on a root TiDB runtime shared by active tests in the current test binary.
func StartInitializedTiDBWithSharedRuntime(t *testing.T, runtime TiDBTestRuntime) string {
	t.Helper()

	rootDSN := sharedTiDBRuntimeRootDSN(t, runtime)
	return StartInitializedTiDBWithRootDSN(t, rootDSN)
}

func sharedTiDBRuntimeRootDSN(t testing.TB, runtime TiDBTestRuntime) string {
	t.Helper()
	if runtime == "" {
		runtime = TiDBTestRuntimeUnistore
	}

	sharedTiDBRuntimeRootMu.Lock()
	for {
		state := sharedTiDBRuntimeRoots[runtime]
		if state == nil {
			state = &sharedTiDBRuntimeRootState{fixture: NewTiDBRuntimeFixture(runtime)}
			sharedTiDBRuntimeRoots[runtime] = state
		}
		if !state.closing {
			state.users++
			sharedTiDBRuntimeRootMu.Unlock()
			t.Cleanup(func() {
				releaseSharedTiDBRuntimeRoot(t, runtime, state)
			})
			return state.fixture.RootDSN(t)
		}
		done := state.done
		sharedTiDBRuntimeRootMu.Unlock()
		<-done
		sharedTiDBRuntimeRootMu.Lock()
	}
}

func releaseSharedTiDBRuntimeRoot(t testing.TB, runtime TiDBTestRuntime, state *sharedTiDBRuntimeRootState) {
	t.Helper()

	var fixture *TiDBRuntimeFixture

	sharedTiDBRuntimeRootMu.Lock()
	if sharedTiDBRuntimeRoots[runtime] != state {
		sharedTiDBRuntimeRootMu.Unlock()
		return
	}
	if state.users > 0 {
		state.users--
	}
	if state.users == 0 && !state.closing {
		state.closing = true
		state.done = make(chan struct{})
		fixture = state.fixture
	}
	sharedTiDBRuntimeRootMu.Unlock()

	if fixture == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	err := fixture.Cleanup(ctx)
	cancel()

	sharedTiDBRuntimeRootMu.Lock()
	if sharedTiDBRuntimeRoots[runtime] == state {
		delete(sharedTiDBRuntimeRoots, runtime)
		close(state.done)
	}
	sharedTiDBRuntimeRootMu.Unlock()

	require.NoError(t, err)
}

// StartInitializedTiDBWithPackageRuntime returns a migrated isolated database
// on a root TiDB runtime retained until RunTestsWithPackageTiDBRuntimeRoots
// cleanup runs.
func StartInitializedTiDBWithPackageRuntime(t *testing.T, runtime TiDBTestRuntime) string {
	t.Helper()

	rootDSN := packageTiDBRuntimeRootDSN(t, runtime)
	return StartInitializedTiDBWithRootDSN(t, rootDSN)
}

func packageTiDBRuntimeRootDSN(t testing.TB, runtime TiDBTestRuntime) string {
	t.Helper()
	if runtime == "" {
		runtime = TiDBTestRuntimeUnistore
	}

	packageTiDBRuntimeRootMu.Lock()
	fixture := packageTiDBRuntimeRoots[runtime]
	if fixture == nil {
		fixture = NewTiDBRuntimeFixture(runtime)
		packageTiDBRuntimeRoots[runtime] = fixture
	}
	packageTiDBRuntimeRootMu.Unlock()

	return fixture.RootDSN(t)
}

func cleanupPackageTiDBRuntimeRoots(ctx context.Context) error {
	packageTiDBRuntimeRootMu.Lock()
	fixtures := make([]*TiDBRuntimeFixture, 0, len(packageTiDBRuntimeRoots))
	for _, fixture := range packageTiDBRuntimeRoots {
		fixtures = append(fixtures, fixture)
	}
	packageTiDBRuntimeRoots = make(map[TiDBTestRuntime]*TiDBRuntimeFixture)
	packageTiDBRuntimeRootMu.Unlock()

	var cleanupErr error
	for _, fixture := range fixtures {
		if fixture == nil {
			continue
		}
		cleanupErr = errors.Join(cleanupErr, fixture.Cleanup(ctx))
	}
	return cleanupErr
}

// RunTestsWithPackageTiDBRuntimeRoots runs a test package and cleans up root
// runtimes retained through StartInitializedTiDBWithPackageRuntime.
func RunTestsWithPackageTiDBRuntimeRoots(m *testing.M, name string) int {
	code := m.Run()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	if err := cleanupPackageTiDBRuntimeRoots(ctx); err != nil {
		fmt.Fprintf(os.Stderr, "cleanup %s tidb runtime roots: %v\n", name, err)
		if code == 0 {
			code = 1
		}
	}
	return code
}

// OpenInitializedTiDBCoreDB opens a migrated TiDB database for package-level
// fixtures and returns an explicit cleanup function using the default test runtime.
func OpenInitializedTiDBCoreDB(ctx context.Context) (coredb.DB, func(context.Context) error, error) {
	return OpenInitializedTiDBCoreDBWithRuntime(ctx, TiDBTestRuntimeUnistore)
}

// OpenInitializedTiDBCoreDBWithRuntime opens a migrated TiDB database for package-level
// fixtures on the requested runtime and returns an explicit cleanup function.
func OpenInitializedTiDBCoreDBWithRuntime(ctx context.Context, runtime TiDBTestRuntime) (coredb.DB, func(context.Context) error, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	dsn, cleanupDatabase, err := StartInitializedTiDBWithRuntimeAndCleanup(ctx, runtime)
	if err != nil {
		return nil, nil, err
	}

	sqlDB, err := OpenTiDBContext(ctx, dsn)
	if err != nil {
		cleanupCtx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer cancel()
		_ = cleanupDatabase(cleanupCtx)
		return nil, nil, err
	}

	cleanup := func(cleanupCtx context.Context) error {
		var cleanupErr error
		if err := sqlDB.Close(); err != nil {
			cleanupErr = err
		}
		if err := cleanupDatabase(cleanupCtx); cleanupErr == nil && err != nil {
			cleanupErr = err
		}
		return cleanupErr
	}
	return coredb.WrapSQLDB(sqlDB), cleanup, nil
}

// BeginTiDBRollbackDB returns a DB scoped to a test transaction that rolls back
// during cleanup. Nested BeginTx calls use TiDB savepoints.
func BeginTiDBRollbackDB(t testing.TB, ctx context.Context, db coredb.DB) coredb.DB {
	t.Helper()
	if ctx == nil {
		ctx = context.Background()
	}

	tx, err := db.BeginTx(ctx, coredb.TxOptions{})
	require.NoError(t, err)
	t.Cleanup(func() {
		if err := tx.Rollback(context.Background()); err != nil {
			t.Errorf("rollback tidb fixture transaction: %v", err)
		}
	})
	return &tiDBRollbackDB{tx: tx}
}

type tiDBRollbackDB struct {
	tx      coredb.Tx
	counter uint64
}

func (db *tiDBRollbackDB) Database() string {
	return "tidb"
}

func (db *tiDBRollbackDB) Exec(ctx context.Context, query string, args ...any) (coredb.Result, error) {
	if db == nil || db.tx == nil {
		return nil, fmt.Errorf("tidb rollback db is nil")
	}
	return db.tx.Exec(ctx, query, args...)
}

func (db *tiDBRollbackDB) Query(ctx context.Context, query string, args ...any) (coredb.Rows, error) {
	if db == nil || db.tx == nil {
		return nil, fmt.Errorf("tidb rollback db is nil")
	}
	return db.tx.Query(ctx, query, args...)
}

func (db *tiDBRollbackDB) QueryRow(ctx context.Context, query string, args ...any) coredb.Row {
	if db == nil || db.tx == nil {
		return errorCoreDBRow{err: fmt.Errorf("tidb rollback db is nil")}
	}
	return db.tx.QueryRow(ctx, query, args...)
}

func (db *tiDBRollbackDB) Get(ctx context.Context, dest any, query string, args ...any) error {
	if db == nil || db.tx == nil {
		return fmt.Errorf("tidb rollback db is nil")
	}
	return coredb.Get(ctx, db.tx, dest, query, args...)
}

func (db *tiDBRollbackDB) Select(ctx context.Context, dest any, query string, args ...any) error {
	if db == nil || db.tx == nil {
		return fmt.Errorf("tidb rollback db is nil")
	}
	return coredb.Select(ctx, db.tx, dest, query, args...)
}

func (db *tiDBRollbackDB) BeginTx(ctx context.Context, options coredb.TxOptions) (coredb.Tx, error) {
	if db == nil || db.tx == nil {
		return nil, fmt.Errorf("tidb rollback db is nil")
	}
	if options != (coredb.TxOptions{}) {
		return nil, fmt.Errorf("tidb rollback db nested transactions do not support tx options")
	}

	savepoint := fmt.Sprintf("agentteam_test_savepoint_%d", atomic.AddUint64(&db.counter, 1))
	if _, err := db.tx.Exec(ctx, "SAVEPOINT "+savepoint); err != nil {
		return nil, fmt.Errorf("create tidb savepoint: %w", err)
	}
	return &tiDBSavepointTx{tx: db.tx, savepoint: savepoint}, nil
}

func (db *tiDBRollbackDB) WithSystemBootstrapLockTx(ctx context.Context, fn func(coredb.Tx) error) (err error) {
	if db == nil || db.tx == nil {
		return fmt.Errorf("tidb rollback db is nil")
	}
	if fn == nil {
		return fmt.Errorf("coredb bootstrap tx function is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	var release func(context.Context) error
	defer func() {
		if release == nil {
			return
		}
		releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
		defer cancel()
		err = errors.Join(err, release(releaseCtx))
	}()

	err = coredb.WithTx(ctx, db, func(tx coredb.Tx) error {
		var lockErr error
		release, lockErr = coredb.AcquireSystemBootstrapLock(ctx, tx)
		if lockErr != nil {
			return lockErr
		}
		return fn(tx)
	})
	return err
}

type tiDBSavepointTx struct {
	tx        coredb.Tx
	savepoint string
	mu        sync.Mutex
	done      bool
}

func (tx *tiDBSavepointTx) Database() string {
	return "tidb"
}

func (tx *tiDBSavepointTx) Exec(ctx context.Context, query string, args ...any) (coredb.Result, error) {
	if tx == nil || tx.tx == nil {
		return nil, fmt.Errorf("tidb savepoint tx is nil")
	}
	return tx.tx.Exec(ctx, query, args...)
}

func (tx *tiDBSavepointTx) Query(ctx context.Context, query string, args ...any) (coredb.Rows, error) {
	if tx == nil || tx.tx == nil {
		return nil, fmt.Errorf("tidb savepoint tx is nil")
	}
	return tx.tx.Query(ctx, query, args...)
}

func (tx *tiDBSavepointTx) QueryRow(ctx context.Context, query string, args ...any) coredb.Row {
	if tx == nil || tx.tx == nil {
		return errorCoreDBRow{err: fmt.Errorf("tidb savepoint tx is nil")}
	}
	return tx.tx.QueryRow(ctx, query, args...)
}

func (tx *tiDBSavepointTx) Get(ctx context.Context, dest any, query string, args ...any) error {
	if tx == nil || tx.tx == nil {
		return fmt.Errorf("tidb savepoint tx is nil")
	}
	return coredb.Get(ctx, tx.tx, dest, query, args...)
}

func (tx *tiDBSavepointTx) Select(ctx context.Context, dest any, query string, args ...any) error {
	if tx == nil || tx.tx == nil {
		return fmt.Errorf("tidb savepoint tx is nil")
	}
	return coredb.Select(ctx, tx.tx, dest, query, args...)
}

func (tx *tiDBSavepointTx) Commit(ctx context.Context) error {
	if tx == nil || tx.tx == nil {
		return fmt.Errorf("tidb savepoint tx is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	return tx.finish(func() error {
		_, err := tx.tx.Exec(ctx, "RELEASE SAVEPOINT "+tx.savepoint)
		return err
	})
}

func (tx *tiDBSavepointTx) Rollback(ctx context.Context) error {
	if tx == nil || tx.tx == nil {
		return fmt.Errorf("tidb savepoint tx is nil")
	}
	if ctx == nil {
		ctx = context.Background()
	}
	return tx.finish(func() error {
		if _, err := tx.tx.Exec(ctx, "ROLLBACK TO SAVEPOINT "+tx.savepoint); err != nil {
			return err
		}
		_, err := tx.tx.Exec(ctx, "RELEASE SAVEPOINT "+tx.savepoint)
		return err
	})
}

func (tx *tiDBSavepointTx) finish(fn func() error) error {
	tx.mu.Lock()
	if tx.done {
		tx.mu.Unlock()
		return fmt.Errorf("tidb savepoint tx already closed")
	}
	tx.done = true
	tx.mu.Unlock()

	return fn()
}

type errorCoreDBRow struct {
	err error
}

func (row errorCoreDBRow) Scan(...any) error {
	return row.err
}