package middleware

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

type tokenBucket struct {
	mu     sync.Mutex
	rate   float64 // tokens per second
	burst  float64 // max tokens
	tokens float64
	last   time.Time
}

func newTokenBucket(max int, per time.Duration, burst int) *tokenBucket {
	if burst <= 0 {
		burst = max
	}
	rate := float64(max) / per.Seconds()
	b := float64(burst)
	now := time.Now()
	return &tokenBucket{rate: rate, burst: b, tokens: b, last: now}
}

func (b *tokenBucket) allow() bool {
	b.mu.Lock()
	defer b.mu.Unlock()

	now := time.Now()
	delta := now.Sub(b.last).Seconds()
	b.last = now

	b.tokens += delta * b.rate
	if b.tokens > b.burst {
		b.tokens = b.burst
	}

	if b.tokens < 1 {
		return false
	}
	b.tokens -= 1
	return true
}

type ipClient struct {
	bucket   *tokenBucket
	max      int
	burst    int
	lastSeen time.Time
}

// RateLimitByIP returns a Gin middleware that rate limits requests per client IP.
//
// max: max requests per time window (per)
// per: the time window duration
// burst: optional burst capacity (defaults to max if <=0)
// ttl: how long to keep idle IP buckets around
func RateLimitByIP(max int, per time.Duration, burst int, ttl time.Duration) gin.HandlerFunc {
	return RateLimitByIPDynamic(func() int { return max }, per, func() int { return burst }, ttl)
}

// RateLimitByIPDynamic is like RateLimitByIP but reads max/burst dynamically.
// This allows changing limits at runtime (e.g. from an admin config page).
func RateLimitByIPDynamic(maxFn func() int, per time.Duration, burstFn func() int, ttl time.Duration) gin.HandlerFunc {
	var (
		mu      sync.Mutex
		clients = make(map[string]*ipClient)
	)

	// opportunistic cleanup (runs at most once per minute)
	var (
		cleanupMu   sync.Mutex
		lastCleanup time.Time
	)

	cleanup := func(now time.Time) {
		cleanupMu.Lock()
		defer cleanupMu.Unlock()
		if !lastCleanup.IsZero() && now.Sub(lastCleanup) < time.Minute {
			return
		}
		lastCleanup = now

		mu.Lock()
		defer mu.Unlock()
		for ip, c := range clients {
			if now.Sub(c.lastSeen) > ttl {
				delete(clients, ip)
			}
		}
	}

	getClient := func(ip string, now time.Time, max int, burst int) *ipClient {
		mu.Lock()
		defer mu.Unlock()
		c, ok := clients[ip]
		if !ok {
			c = &ipClient{bucket: newTokenBucket(max, per, burst), max: max, burst: burst, lastSeen: now}
			clients[ip] = c
			return c
		}
		c.lastSeen = now
		// If settings changed, reset the bucket.
		if c.max != max || c.burst != burst {
			c.bucket = newTokenBucket(max, per, burst)
			c.max = max
			c.burst = burst
		}
		return c
	}

	return func(c *gin.Context) {
		now := time.Now()
		cleanup(now)

		ip := c.ClientIP()
		max := maxFn()
		if max <= 0 {
			max = 1
		}
		burst := burstFn()
		if burst <= 0 {
			burst = max
		}
		client := getClient(ip, now, max, burst)

		if !client.bucket.allow() {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error": "rate limit exceeded",
			})
			return
		}

		c.Next()
	}
}