package middleware

import (
	"sync"
	"time"

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

// TokenBucket implements a token bucket rate limiter per key (e.g. IP or device ID).
type TokenBucket struct {
	mu       sync.Mutex
	buckets  map[string]*bucket
	rate     float64 // tokens per second
	capacity float64
}

type bucket struct {
	tokens   float64
	lastFill time.Time
}

func NewTokenBucket(rate, capacity int) *TokenBucket {
	return &TokenBucket{
		buckets:  make(map[string]*bucket),
		rate:     float64(rate),
		capacity: float64(capacity),
	}
}

func (tb *TokenBucket) Allow(key string) bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	b, ok := tb.buckets[key]
	if !ok {
		b = &bucket{tokens: tb.capacity, lastFill: time.Now()}
		tb.buckets[key] = b
	}

	now := time.Now()
	elapsed := now.Sub(b.lastFill).Seconds()
	b.tokens += elapsed * tb.rate
	if b.tokens > tb.capacity {
		b.tokens = tb.capacity
	}
	b.lastFill = now

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

// Cleanup removes stale buckets older than the given duration.
func (tb *TokenBucket) Cleanup(maxAge time.Duration) {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	for k, b := range tb.buckets {
		if now.Sub(b.lastFill) > maxAge {
			delete(tb.buckets, k)
		}
	}
}

// RateLimitMiddleware returns a Gin middleware that rate-limits requests per client IP.
func RateLimitMiddleware(tb *TokenBucket) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		if !tb.Allow(ip) {
			c.AbortWithStatusJSON(429, gin.H{"error": "Rate limit exceeded. Please slow down."})
			return
		}
		c.Next()
	}
}

// GatedSemaphore limits the number of concurrent operations (e.g. downloads).
type GatedSemaphore struct {
	ch chan struct{}
}

func NewGatedSemaphore(maxConcurrent int) *GatedSemaphore {
	return &GatedSemaphore{ch: make(chan struct{}, maxConcurrent)}
}

func (gs *GatedSemaphore) Acquire() {
	gs.ch <- struct{}{}
}

func (gs *GatedSemaphore) Release() {
	<-gs.ch
}

func (gs *GatedSemaphore) Available() int {
	return cap(gs.ch) - len(gs.ch)
}

// GatedMiddleware only allows requests through if the semaphore has capacity.
func GatedMiddleware(gs *GatedSemaphore) gin.HandlerFunc {
	return func(c *gin.Context) {
		select {
		case gs.ch <- struct{}{}:
			defer func() { <-gs.ch }()
			c.Next()
		default:
			c.AbortWithStatusJSON(503, gin.H{"error": "Server at capacity. Please retry later."})
		}
	}
}