package middleware

import (
	"container/list"
	"errors"
	"math"
	"strconv"
	"sync"
	"time"

	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
)

type RateLimitConfig struct {
	Rate    float64                     `json:"rate"` // Tokens per second, per key.
	Burst   int                         `json:"burst"`
	MaxKeys int                         `json:"maxKeys"`
	IdleTTL time.Duration               `json:"idleTtl"`
	Key     func(*httpx.Context) string `json:"-"` // Default: trusted ClientIP.
}

type bucket struct {
	key     string
	tokens  float64
	updated time.Time
}
type RateLimiter struct {
	mu     sync.Mutex
	config RateLimitConfig
	items  map[string]*list.Element
	lru    *list.List
	now    func() time.Time
}

func NewRateLimiter(cfg RateLimitConfig) (*RateLimiter, error) {
	if math.IsNaN(cfg.Rate) || math.IsInf(cfg.Rate, 0) || cfg.Rate <= 0 || cfg.Burst < 1 || cfg.MaxKeys < 0 || cfg.IdleTTL < 0 {
		return nil, errors.New("middleware: invalid rate limit configuration")
	}
	refill := float64(cfg.Burst) / cfg.Rate
	if refill > (24 * time.Hour).Seconds() {
		return nil, errors.New("middleware: bucket refill time must not exceed 24h")
	}
	minimum := time.Duration(math.Ceil(refill * float64(time.Second)))
	if cfg.IdleTTL == 0 {
		cfg.IdleTTL = max(5*time.Minute, minimum)
	}
	if cfg.IdleTTL < minimum {
		return nil, errors.New("middleware: idleTtl must allow a full bucket refill")
	}
	if cfg.MaxKeys == 0 {
		cfg.MaxKeys = 10000
	}
	if cfg.Key == nil {
		cfg.Key = func(c *httpx.Context) string { return c.ClientIP() }
	}
	return &RateLimiter{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil
}

func (l *RateLimiter) Middleware() httpx.Middleware {
	return func(next httpx.Handler) httpx.Handler {
		return func(c *httpx.Context) error {
			key := l.config.Key(c)
			if key == "" || len(key) > 1024 {
				return apierrors.BadRequest("INVALID_RATE_KEY", "invalid rate limit key")
			}
			delay, full := l.allow(key)
			if full {
				return apierrors.New(503, "RATE_LIMIT_CAPACITY", "rate limiter capacity reached")
			}
			if delay > 0 {
				c.SetResponseHeader("Retry-After", strconv.Itoa(max(1, int(math.Ceil(delay.Seconds())))))
				return apierrors.New(429, "RATE_LIMITED", "too many requests")
			}
			return next(c)
		}
	}
}

func (l *RateLimiter) allow(key string) (time.Duration, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()
	e := l.items[key]
	if e == nil {
		for l.lru.Len() > 0 {
			old := l.lru.Back()
			if now.Sub(old.Value.(*bucket).updated) < l.config.IdleTTL {
				break
			}
			delete(l.items, old.Value.(*bucket).key)
			l.lru.Remove(old)
		}
		if len(l.items) >= l.config.MaxKeys {
			return 0, true
		}
		e = l.lru.PushFront(&bucket{key: key, tokens: float64(l.config.Burst), updated: now})
		l.items[key] = e
	}
	b := e.Value.(*bucket)
	b.tokens = min(float64(l.config.Burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*l.config.Rate)
	b.updated = now
	l.lru.MoveToFront(e)
	if b.tokens >= 1 {
		b.tokens--
		return 0, false
	}
	return time.Duration(math.Ceil((1 - b.tokens) / l.config.Rate * float64(time.Second))), false
}