package billing

import (
	"fmt"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/go-logr/logr"
	"github.com/tidbcloud/devops/api"
	ctrl "github.com/tidbcloud/devops/pkg/controller"
	"github.com/tidbcloud/devops/pkg/infra/config"
	"github.com/tidbcloud/devops/pkg/infra/consts"
	"golang.org/x/time/rate"
)

const maxLimiterEntries = 4096

type proxyRateLimiter struct {
	cfg    config.BillingProxyRateLimitConfig
	logger logr.Logger

	mu      sync.Mutex
	entries map[string]*rate.Limiter
}

func newProxyRateLimiter(cfg config.BillingProxyRateLimitConfig, logger logr.Logger) *proxyRateLimiter {
	return &proxyRateLimiter{
		cfg:     cfg,
		logger:  logger,
		entries: make(map[string]*rate.Limiter),
	}
}

func (l *proxyRateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !l.cfg.Enabled {
			c.Next()
			return
		}

		operatorType := c.GetString(consts.CtxOperatorType)
		operator := c.GetString(consts.CtxOperator)
		key := operatorType + ":" + operator
		if !l.allow(key, operatorType) {
			l.logger.Info("billing proxy rate limited",
				"operator_type", operatorType,
				"operator", operator,
				"path", c.Request.URL.Path,
			)
			ctrl.FailWithCode(c, fmt.Errorf("rate limit exceeded, retry later"), api.ErrCodeResourceExhaust)
			return
		}
		c.Next()
	}
}

func (l *proxyRateLimiter) allow(key, operatorType string) bool {
	limiter := l.getLimiter(key, operatorType)
	return limiter.Allow()
}

func (l *proxyRateLimiter) getLimiter(key, operatorType string) *rate.Limiter {
	l.mu.Lock()
	defer l.mu.Unlock()

	if limiter, ok := l.entries[key]; ok {
		return limiter
	}

	if len(l.entries) >= maxLimiterEntries {
		l.logger.Info("billing proxy rate limiter map full, clearing entries", "entries", len(l.entries))
		l.entries = make(map[string]*rate.Limiter)
	}

	rps, burst := limitsForOperatorType(l.cfg, operatorType)
	limiter := rate.NewLimiter(rate.Limit(rps), burst)
	l.entries[key] = limiter
	return limiter
}

func limitsForOperatorType(cfg config.BillingProxyRateLimitConfig, operatorType string) (float64, int) {
	switch operatorType {
	case consts.CtxOperatorTypeAPIToken:
		return bucketLimits(cfg.APIToken, 0.5, 10)
	case consts.CtxOperatorTypeAPIKey:
		return bucketLimits(cfg.APIKey, 0.5, 10)
	default:
		return bucketLimits(cfg.User, 2, 30)
	}
}

func bucketLimits(cfg config.RateLimitBucketConfig, defaultRPS float64, defaultBurst int) (float64, int) {
	rps := cfg.RPS
	if rps <= 0 {
		rps = defaultRPS
	}
	burst := cfg.Burst
	if burst <= 0 {
		burst = defaultBurst
	}
	return rps, burst
}