package sqltx

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"strings"
	"time"
)

const pessimisticTxCleanupTimeout = 200 * time.Millisecond

const (
	tidbConstraintCheckInPlaceSelect = "SELECT @@SESSION.tidb_constraint_check_in_place_pessimistic"
	tidbConstraintCheckInPlaceOn     = "SET SESSION tidb_constraint_check_in_place_pessimistic = ON"
)

// BeginPessimisticTx acquires a dedicated *sql.Conn, sets SESSION-level TiDB
// pessimistic mode on that conn, and begins a transaction. The cleanup closure
// must be called after commit or rollback so the session variable is reset
// before returning the connection to the pool.
func BeginPessimisticTx(ctx context.Context, db *sql.DB) (*sql.Tx, func() error, error) {
	return beginPessimisticTx(ctx, db, nil, false)
}

// BeginPessimisticConstraintCheckingTx starts a pessimistic transaction whose
// uniqueness and foreign-key constraints are checked by each statement. It is
// intended for transactions that must classify a constraint conflict before
// Commit. The cleanup closure restores the session's prior constraint-check
// value as well as tidb_txn_mode before returning the connection to the pool.
func BeginPessimisticConstraintCheckingTx(ctx context.Context, db *sql.DB) (*sql.Tx, func() error, error) {
	return beginPessimisticTx(ctx, db, nil, true)
}

// BeginPessimisticReadCommittedTx starts a pessimistic transaction whose
// snapshot is refreshed for each statement. Use it when a transaction must
// serialize writes on one row while observing rows committed by a waiter.
func BeginPessimisticReadCommittedTx(ctx context.Context, db *sql.DB) (*sql.Tx, func() error, error) {
	return beginPessimisticTx(ctx, db, &sql.TxOptions{Isolation: sql.LevelReadCommitted}, false)
}

func beginPessimisticTx(ctx context.Context, db *sql.DB, opts *sql.TxOptions, checkConstraintsInPlace bool) (*sql.Tx, func() error, error) {
	const maxOpenAttempts = 8

	var lastErr error
	for attempt := 1; attempt <= maxOpenAttempts; attempt++ {
		tx, cleanup, err := beginPessimisticTxOnce(ctx, db, opts, checkConstraintsInPlace)
		if err == nil {
			return tx, cleanup, nil
		}
		lastErr = err
		if ctx.Err() != nil || !isRetryablePessimisticOpenErr(err) {
			return nil, nil, err
		}
		if attempt < maxOpenAttempts {
			delay := time.Duration(attempt*25) * time.Millisecond
			timer := time.NewTimer(delay)
			select {
			case <-ctx.Done():
				timer.Stop()
				return nil, nil, ctx.Err()
			case <-timer.C:
			}
		}
	}
	return nil, nil, fmt.Errorf("begin pessimistic tx after %d attempts: %w", maxOpenAttempts, lastErr)
}

func beginPessimisticTxOnce(ctx context.Context, db *sql.DB, opts *sql.TxOptions, checkConstraintsInPlace bool) (*sql.Tx, func() error, error) {
	conn, err := db.Conn(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("acquire conn: %w", err)
	}

	var previousConstraintCheck string
	if checkConstraintsInPlace {
		if err := conn.QueryRowContext(ctx, tidbConstraintCheckInPlaceSelect).Scan(&previousConstraintCheck); err != nil {
			_ = conn.Close()
			return nil, nil, fmt.Errorf("read tidb_constraint_check_in_place_pessimistic: %w", err)
		}
		previousConstraintCheck, err = canonicalTiDBBoolean(previousConstraintCheck)
		if err != nil {
			_ = conn.Close()
			return nil, nil, err
		}
	}

	cleanup := func() error {
		cleanupCtx, cancel := context.WithTimeout(context.Background(), pessimisticTxCleanupTimeout)
		defer cancel()

		if checkConstraintsInPlace {
			restoreConstraintErr := restoreConstraintCheck(cleanupCtx, conn, previousConstraintCheck)
			resetTxnModeErr := resetTxnModeDefault(cleanupCtx, conn)
			resetErr := errors.Join(restoreConstraintErr, resetTxnModeErr)
			if resetErr != nil {
				markConnBad(conn)
			}
			return errors.Join(resetErr, conn.Close())
		}

		resetErr := resetTxnModeDefault(cleanupCtx, conn)
		if resetErr != nil {
			markConnBad(conn)
		}
		closeErr := conn.Close()
		if resetErr != nil {
			return resetErr
		}
		return closeErr
	}

	if _, err := conn.ExecContext(ctx, "SET SESSION tidb_txn_mode = 'pessimistic'"); err != nil {
		_ = cleanup()
		return nil, nil, fmt.Errorf("set tidb_txn_mode=pessimistic: %w", err)
	}
	if checkConstraintsInPlace {
		if _, err := conn.ExecContext(ctx, tidbConstraintCheckInPlaceOn); err != nil {
			_ = cleanup()
			return nil, nil, fmt.Errorf("set tidb_constraint_check_in_place_pessimistic=ON: %w", err)
		}
	}

	tx, err := conn.BeginTx(ctx, opts)
	if err != nil {
		_ = cleanup()
		return nil, nil, fmt.Errorf("begin tx: %w", err)
	}
	return tx, cleanup, nil
}

func canonicalTiDBBoolean(value string) (string, error) {
	switch strings.ToUpper(strings.TrimSpace(value)) {
	case "ON", "1":
		return "ON", nil
	case "OFF", "0":
		return "OFF", nil
	default:
		return "", fmt.Errorf("unexpected tidb_constraint_check_in_place_pessimistic value %q", value)
	}
}

func restoreConstraintCheck(ctx context.Context, conn *sql.Conn, previous string) error {
	if _, err := conn.ExecContext(ctx, "SET SESSION tidb_constraint_check_in_place_pessimistic = "+previous); err != nil {
		return fmt.Errorf("restore tidb_constraint_check_in_place_pessimistic=%s: %w", previous, err)
	}
	return nil
}

func isRetryablePessimisticOpenErr(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, driver.ErrBadConn) || errors.Is(err, sql.ErrConnDone) {
		return true
	}
	msg := strings.ToLower(err.Error())
	if strings.Contains(msg, "driver: bad connection") || strings.Contains(msg, "bad connection") {
		return true
	}
	return false
}

func resetTxnModeDefault(ctx context.Context, conn *sql.Conn) error {
	if _, err := conn.ExecContext(ctx, "SET SESSION tidb_txn_mode = DEFAULT"); err != nil {
		return fmt.Errorf("reset tidb_txn_mode to DEFAULT: %w", err)
	}
	return nil
}

func markConnBad(conn *sql.Conn) {
	_ = conn.Raw(func(any) error {
		return driver.ErrBadConn
	})
}