package tidbmanager

import (
	"context"
	"strconv"
	"sync"

	"github.com/go-logr/logr"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"golang.org/x/time/rate"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const rateLimiterBurst = 1

type rateLimiter struct {
	logr.Logger
	client            client.Client
	tidbPoolNamespace string

	limiterMu sync.Mutex
	limiter   *rate.Limiter
	rateLimit int
}

func newRateLimiter(
	logger logr.Logger,
	tidbPoolNamespace string,
	kubeClient client.Client,
) *rateLimiter {
	return &rateLimiter{
		Logger:            logger.WithName("rateLimiter"),
		client:            kubeClient,
		tidbPoolNamespace: tidbPoolNamespace,
	}
}

func (r *rateLimiter) wait(ctx context.Context) error {
	if r == nil {
		return nil
	}

	limit := r.getRateLimit(ctx)
	if limit <= 0 {
		r.limiterMu.Lock()
		r.limiter = nil
		r.rateLimit = 0
		r.limiterMu.Unlock()
		return nil
	}

	r.limiterMu.Lock()
	if r.limiter == nil || r.rateLimit != limit {
		r.limiter = rate.NewLimiter(rate.Limit(limit), rateLimiterBurst)
		r.rateLimit = limit
	}
	limiter := r.limiter
	r.limiterMu.Unlock()

	return limiter.Wait(ctx)
}

func (r *rateLimiter) getRateLimit(ctx context.Context) int {
	serverlessPoolConfigCm := &corev1.ConfigMap{}
	if err := r.client.Get(ctx, client.ObjectKey{
		Name:      types.ServerlessPoolConfig,
		Namespace: r.tidbPoolNamespace,
	}, serverlessPoolConfigCm); err != nil {
		r.Error(err, "get cm serverlessPoolConfig for rateLimiter")
		return 0
	}

	limit, err := strconv.Atoi(serverlessPoolConfigCm.Data[types.AssignRateLimitPerManager])
	if err != nil || limit <= 0 {
		return 0
	}
	return limit
}