package action

import (
	"context"
	"errors"
	"time"

	"github.com/shiningrush/fastflow/pkg/entity/run"
	pkglog "github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

var ErrStopRetry = errors.New("stop retry")

type RetryFun func(ctx context.Context, logger *zap.Logger, run func() error) error

type RunAction interface {
	Params() any
	Run(ctx context.Context, logger *zap.Logger, params any) error
}

type RunAfterAction interface {
	RunAfter(ctx context.Context, logger *zap.Logger, params any) error
}

type BaseActionOpt func(config *BaseActionConfig)

type BaseActionConfig struct {
	runRetryFun      RetryFun
	runAfterRetryFun RetryFun
}

type Wrapper struct {
	name     string
	run      RunAction
	runAfter RunAfterAction
	config   *BaseActionConfig
}

func MustNewActionWrapper(name string, action any, opts ...BaseActionOpt) *Wrapper {
	runAction, ok := action.(RunAction)
	if !ok {
		panic("action must implement RunAction interface")
	}

	wrapper := &Wrapper{
		name: name,
		run:  runAction,
		config: &BaseActionConfig{
			runRetryFun:      InvokeRunActionWithExponentialRetry,
			runAfterRetryFun: InvokeRunActionWithInfinitelyRetry,
		},
	}
	if runAfterAction, ok := action.(RunAfterAction); ok {
		wrapper.runAfter = runAfterAction
	}
	for _, opt := range opts {
		opt(wrapper.config)
	}
	return wrapper
}

func (m *Wrapper) Name() string {
	return m.name
}

func (m *Wrapper) ParameterNew() any {
	return m.run.Params()
}

func (m *Wrapper) Run(eCtx run.ExecuteContext, params any) error {
	logger := pkglog.WithContextL(eCtx.Context()).Named(m.Name())
	return m.config.runRetryFun(eCtx.Context(), logger, func() error {
		return m.run.Run(eCtx.Context(), logger, params)
	})
}

func (m *Wrapper) RunAfter(eCtx run.ExecuteContext, params any) error {
	if m.runAfter == nil {
		return nil
	}
	logger := pkglog.WithContextL(eCtx.Context()).Named(m.Name())
	return m.config.runAfterRetryFun(eCtx.Context(), logger, func() error {
		return m.runAfter.RunAfter(eCtx.Context(), logger, params)
	})
}

func SetRunRetryFun(retryFun RetryFun) BaseActionOpt {
	return func(config *BaseActionConfig) {
		config.runRetryFun = retryFun
	}
}

func SetRunAfterRetryFun(retryFun RetryFun) BaseActionOpt {
	return func(config *BaseActionConfig) {
		config.runAfterRetryFun = retryFun
	}
}

func InvokeRunActionWithExponentialRetry(ctx context.Context, logger *zap.Logger, run func() error) error {
	return invokeWithRetry(ctx, logger, run, retryConfig{
		delay:       500 * time.Millisecond,
		maxDelay:    10 * time.Second,
		attempts:    8,
		exponential: true,
	})
}

func InvokeRunActionWithInfinitelyRetry(ctx context.Context, logger *zap.Logger, run func() error) error {
	return invokeWithRetry(ctx, logger, run, retryConfig{
		delay:    15 * time.Second,
		attempts: 0,
	})
}

type retryConfig struct {
	delay       time.Duration
	maxDelay    time.Duration
	attempts    int
	exponential bool
}

func invokeWithRetry(ctx context.Context, logger *zap.Logger, run func() error, cfg retryConfig) error {
	delay := cfg.delay
	var lastErr error
	for attempt := 0; ; attempt++ {
		err := run()
		if err == nil {
			return nil
		}
		if errors.Is(err, ErrStopRetry) {
			return err
		}
		lastErr = err
		if cfg.attempts > 0 && attempt+1 >= cfg.attempts {
			return lastErr
		}
		logger.Warn("action run error", zap.Int("retry_times", attempt), zap.Error(err))
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return ctx.Err()
		case <-timer.C:
		}
		if cfg.exponential && delay < cfg.maxDelay {
			delay *= 2
			if delay > cfg.maxDelay {
				delay = cfg.maxDelay
			}
		}
	}
}