package executor

import (
	"context"
	"database/sql"
	"math/rand"
	"time"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/cqlerrors"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/metrics"
)

// deadlockRetryConfig controls transparent server-side retry of write
// statements that failed with a deadlock (MySQL 1213). A deadlock fully rolls
// back the transaction, so re-running the whole unit is safe.
type deadlockRetryConfig struct {
	enabled     bool
	maxAttempts int
	baseBackoff time.Duration
	maxBackoff  time.Duration
}

// Defaults for deadlock retry when the operator leaves settings unset.
const (
	defaultDeadlockRetryMaxAttempts = 3
	defaultDeadlockRetryBaseBackoff = 10 * time.Millisecond
	defaultDeadlockRetryMaxBackoff  = 100 * time.Millisecond
)

// newDeadlockRetryConfig resolves the effective deadlock retry settings from raw
// config values, applying defaults. Retry is enabled unless explicitly disabled
// (enabled == nil means on).
func newDeadlockRetryConfig(enabled *bool, maxAttempts int, base, max time.Duration) deadlockRetryConfig {
	cfg := deadlockRetryConfig{
		enabled:     enabled == nil || *enabled,
		maxAttempts: maxAttempts,
		baseBackoff: base,
		maxBackoff:  max,
	}
	if cfg.maxAttempts <= 0 {
		cfg.maxAttempts = defaultDeadlockRetryMaxAttempts
	}
	if cfg.baseBackoff <= 0 {
		cfg.baseBackoff = defaultDeadlockRetryBaseBackoff
	}
	if cfg.maxBackoff <= 0 {
		cfg.maxBackoff = defaultDeadlockRetryMaxBackoff
	}
	if cfg.maxBackoff < cfg.baseBackoff {
		cfg.maxBackoff = cfg.baseBackoff
	}
	return cfg
}

// backoffWaiter sleeps for a duration bounded by ceiling, returning a non-nil
// error if the wait was aborted (e.g. context cancellation).
type backoffWaiter func(ctx context.Context, ceiling time.Duration) error

// backoffCeiling returns the maximum backoff for a zero-based retry index: base
// doubled per attempt, clamped to max.
func backoffCeiling(attempt int, base, max time.Duration) time.Duration {
	if base <= 0 {
		return 0
	}
	ceiling := base
	for i := 0; i < attempt; i++ {
		ceiling *= 2
		if ceiling >= max {
			return max
		}
	}
	if ceiling > max {
		return max
	}
	return ceiling
}

// retryOnDeadlock runs fn, transparently retrying when it returns a
// deadlock-classified error, up to cfg.maxAttempts total attempts. Between
// attempts it calls wait with the backoff ceiling for that retry. It returns the
// final result, the number of attempts made, and the final error.
func retryOnDeadlock(
	ctx context.Context,
	cfg deadlockRetryConfig,
	wait backoffWaiter,
	fn func() (*QueryResult, error),
) (*QueryResult, int, error) {
	maxAttempts := cfg.maxAttempts
	if !cfg.enabled || maxAttempts < 1 {
		maxAttempts = 1
	}

	var (
		result *QueryResult
		err    error
	)
	for attempt := 0; attempt < maxAttempts; attempt++ {
		result, err = fn()
		if err == nil {
			return result, attempt + 1, nil
		}
		if cqlerrors.RetryKindOf(err) != cqlerrors.RetryDeadlock {
			return result, attempt + 1, err
		}
		if attempt == maxAttempts-1 {
			return result, attempt + 1, err
		}
		if waitErr := wait(ctx, backoffCeiling(attempt, cfg.baseBackoff, cfg.maxBackoff)); waitErr != nil {
			return result, attempt + 1, err
		}
	}
	return result, maxAttempts, err
}

// batchContextKey marks a context as executing a batch request.
type batchContextKey struct{}

// markBatchContext flags ctx as a batch request; statements run under it skip
// deadlock retry (batch requests are not retried).
func markBatchContext(ctx context.Context) context.Context {
	return context.WithValue(ctx, batchContextKey{}, true)
}

// isBatchContext reports whether ctx was flagged as a batch request.
func isBatchContext(ctx context.Context) bool {
	flagged, _ := ctx.Value(batchContextKey{}).(bool)
	return flagged
}

// executeWriteWithRetry runs fn — a single attempt at a logical write executed
// via exec — with transparent deadlock retry. Retry is limited to standalone
// single statements: it is skipped inside a caller-owned transaction (*sql.Tx)
// and for any statement that is part of a batch request, since re-running one
// statement of a multi-statement unit would be incorrect.
func (e *Executor) executeWriteWithRetry(ctx context.Context, exec sqlExecutor, fn func() (*QueryResult, error)) (*QueryResult, error) {
	if _, isTx := exec.(*sql.Tx); isTx || isBatchContext(ctx) {
		return fn()
	}
	result, attempts, err := retryOnDeadlock(ctx, e.deadlockRetry, deadlockBackoffWait, fn)
	e.observeDeadlockRetry(ctx, attempts, err)
	return result, err
}

// deadlockRetryOutcome maps an attempt count and final error to a retry-metric
// outcome. It reports record=false when no retry occurred (a single attempt).
func deadlockRetryOutcome(attempts int, err error) (string, bool) {
	if attempts <= 1 {
		return "", false
	}
	if err != nil {
		return metrics.WriteRetryOutcomeExhausted, true
	}
	return metrics.WriteRetryOutcomeSucceeded, true
}

// observeDeadlockRetry records a retry-outcome metric when one or more retries
// were made.
func (e *Executor) observeDeadlockRetry(ctx context.Context, attempts int, err error) {
	if outcome, record := deadlockRetryOutcome(attempts, err); record {
		metrics.RecordWriteRetry(ctx, metrics.WriteRetryKindDeadlock, outcome)
	}
}

// deadlockBackoffWait is the production backoffWaiter: it sleeps for a
// full-jitter duration in [0, ceiling], aborting early if ctx is done.
func deadlockBackoffWait(ctx context.Context, ceiling time.Duration) error {
	if ceiling <= 0 {
		return ctx.Err()
	}
	d := time.Duration(rand.Int63n(int64(ceiling) + 1))
	timer := time.NewTimer(d)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}