package persistence

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

// lockWaitTimeoutCtxKey is the context key for the per-session
// innodb_lock_wait_timeout override (in seconds).
type lockWaitTimeoutCtxKey struct{}

// readCommittedCtxKey is the context key that opts a pessimistic transaction
// into SESSION-level READ COMMITTED isolation.
type readCommittedCtxKey struct{}

// readCommittedPessimisticTxMarkerKey records the isolation that was actually
// selected before BeginTx. It is distinct from readCommittedCtxKey: callers
// may request READ COMMITTED on a context that is already inside an unrelated
// transaction, but permission authoring must verify the active transaction.
type readCommittedPessimisticTxMarkerKey struct{}

// WithReadCommittedIsolation returns a context that asks BeginPessimisticTx to
// run the pinned pessimistic connection at READ COMMITTED instead of the TiDB
// default REPEATABLE READ.
//
// Under REPEATABLE READ a plain SELECT reads the transaction-start snapshot, so
// rollup recompute must use a locking current-read (SELECT ... FOR UPDATE) to
// observe sibling rows committed after the transaction began. Under READ
// COMMITTED each statement reads the latest committed data, so a recompute whose
// parent row is already locked can read with a plain SELECT and stay correct
// without locking every child row. Callers that set this flag MUST hold the
// parent-row lock before the recompute so serialization is preserved.
func WithReadCommittedIsolation(ctx context.Context) context.Context {
	return context.WithValue(ctx, readCommittedCtxKey{}, true)
}

// readCommittedIsolationFromContext reports whether ctx opted into READ
// COMMITTED isolation for the pessimistic connection.
func readCommittedIsolationFromContext(ctx context.Context) bool {
	v, ok := ctx.Value(readCommittedCtxKey{}).(bool)
	return ok && v
}

// withLockWaitTimeoutSecs returns a context that carries the given lock-wait
// timeout (seconds). Values ≤ 0 are ignored and the original context is
// returned unchanged.
func withLockWaitTimeoutSecs(ctx context.Context, secs int) context.Context {
	if secs <= 0 {
		return ctx
	}
	return context.WithValue(ctx, lockWaitTimeoutCtxKey{}, secs)
}

// lockWaitTimeoutSecsFromContext returns the lock-wait timeout stored in ctx
// and true, or 0 and false when the value is absent or non-positive.
func lockWaitTimeoutSecsFromContext(ctx context.Context) (int, bool) {
	v, ok := ctx.Value(lockWaitTimeoutCtxKey{}).(int)
	if !ok || v <= 0 {
		return 0, false
	}
	return v, true
}

// BeginPessimisticTx acquires a dedicated *sql.Conn, sets SESSION-level
// TiDB pessimistic mode on that conn, and begins a transaction. It returns
// the *sql.Tx plus a cleanup closure that the caller MUST invoke after
// commit/rollback.
//
// Pinned *sql.Conn is required: without it, db.Begin() may pull a fresh
// pooled conn where SET SESSION was never applied (tidb.go DSN does not
// set tidb_txn_mode). Cleanup resets the session var before returning the
// conn to the pool; otherwise later unrelated callers can inherit
// pessimistic mode.
//
// If ctx carries a lock-wait timeout value (set via withLockWaitTimeoutSecs),
// this function additionally sets innodb_lock_wait_timeout for the session and
// resets it to DEFAULT in cleanup. The timeout is clamped to [1, 3600] seconds.
func BeginPessimisticTx(ctx context.Context, db *sql.DB) (*sql.Tx, func() error, error) {
	// Pinned *sql.Conn does not get the same transparent stale-connection
	// retry that db.Exec/db.Begin paths usually receive. Test runners can have
	// a pocket of stale TiDB Cloud conns after deploy/restart, so give the pool
	// enough attempts to drain them without hiding non-connection failures.
	const maxOpenAttempts = 8

	var lastErr error
	for attempt := 1; attempt <= maxOpenAttempts; attempt++ {
		tx, cleanup, err := beginPessimisticTxOnce(ctx, db)
		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) (*sql.Tx, func() error, error) {
	conn, err := db.Conn(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("acquire conn: %w", err)
	}

	// lockWaitSet / isolationSet track whether the matching session var was
	// applied so cleanup knows whether to reset it.
	lockWaitSet := false
	isolationSet := false

	cleanup := func() error {
		bgCtx := context.Background()
		// Reset all session vars and ALWAYS close the conn. A reset failure must
		// not skip Close (that would leak the connection); record the first error,
		// mark the conn bad, and still proceed to reset the other vars + Close.
		var resetErr error
		if lockWaitSet {
			if _, err := conn.ExecContext(bgCtx, "SET SESSION innodb_lock_wait_timeout = DEFAULT"); err != nil {
				resetErr = fmt.Errorf("reset innodb_lock_wait_timeout to DEFAULT: %w", err)
			}
		}
		if isolationSet {
			if _, err := conn.ExecContext(bgCtx, "SET SESSION transaction_isolation = DEFAULT"); err != nil && resetErr == nil {
				resetErr = fmt.Errorf("reset transaction_isolation to DEFAULT: %w", err)
			}
		}
		if err := resetTxnModeDefault(bgCtx, conn); err != nil && resetErr == nil {
			resetErr = err
		}
		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)
	}

	// Apply per-session lock-wait timeout if the context carries one.
	if secs, ok := lockWaitTimeoutSecsFromContext(ctx); ok {
		// Clamp to a sane range.
		if secs < 1 {
			secs = 1
		}
		if secs > 3600 {
			secs = 3600
		}
		if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION innodb_lock_wait_timeout = %d", secs)); err != nil {
			_ = cleanup()
			return nil, nil, fmt.Errorf("set innodb_lock_wait_timeout=%d: %w", secs, err)
		}
		lockWaitSet = true
	}

	// Apply READ COMMITTED isolation if the context opted in. This must be set
	// before BeginTx so the transaction runs at RC. The matching reset lives in
	// cleanup, gated by isolationSet.
	if readCommittedIsolationFromContext(ctx) {
		if _, err := conn.ExecContext(ctx, "SET SESSION transaction_isolation = 'READ-COMMITTED'"); err != nil {
			_ = cleanup()
			return nil, nil, fmt.Errorf("set transaction_isolation=READ-COMMITTED: %w", err)
		}
		isolationSet = true
	}

	tx, err := conn.BeginTx(ctx, nil)
	if err != nil {
		_ = cleanup()
		return nil, nil, fmt.Errorf("begin tx: %w", err)
	}
	return tx, cleanup, 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
	})
}