package ratelimit

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

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

type client struct {
	tokens     int
	lastRefill time.Time
}

type RateLimiter struct {
	clients       map[string]*client
	mu            sync.RWMutex
	maxTokens     int
	refillRate    time.Duration
	tokensPerHit  int
	cleanupTicker *time.Ticker
}

// NewRateLimiter creates a new rate limiter
// maxTokens: maximum number of tokens a client can have
// refillRate: how often to add tokens back
// tokensPerHit: tokens consumed per request
func NewRateLimiter(maxTokens int, refillRate time.Duration, tokensPerHit int) *RateLimiter {
	rl := &RateLimiter{
		clients:      make(map[string]*client),
		maxTokens:    maxTokens,
		refillRate:   refillRate,
		tokensPerHit: tokensPerHit,
	}

	// Cleanup old clients every 5 minutes
	rl.cleanupTicker = time.NewTicker(5 * time.Minute)
	go rl.cleanupClients()

	return rl
}

func (rl *RateLimiter) getClient(ip string) *client {
	rl.mu.RLock()
	c, exists := rl.clients[ip]
	rl.mu.RUnlock()

	if !exists {
		rl.mu.Lock()
		c = &client{
			tokens:     rl.maxTokens,
			lastRefill: time.Now(),
		}
		rl.clients[ip] = c
		rl.mu.Unlock()
	}

	return c
}

func (rl *RateLimiter) refillTokens(c *client) {
	now := time.Now()
	elapsed := now.Sub(c.lastRefill)
	tokensToAdd := int(elapsed / rl.refillRate)

	if tokensToAdd > 0 {
		c.tokens += tokensToAdd
		if c.tokens > rl.maxTokens {
			c.tokens = rl.maxTokens
		}
		c.lastRefill = now
	}
}

func (rl *RateLimiter) cleanupClients() {
	for range rl.cleanupTicker.C {
		rl.mu.Lock()
		now := time.Now()
		for ip, c := range rl.clients {
			if now.Sub(c.lastRefill) > 10*time.Minute {
				delete(rl.clients, ip)
			}
		}
		rl.mu.Unlock()
	}
}

// Middleware returns a Gin middleware handler
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		client := rl.getClient(ip)

		rl.mu.Lock()
		rl.refillTokens(client)

		if client.tokens >= rl.tokensPerHit {
			client.tokens -= rl.tokensPerHit
			rl.mu.Unlock()
			c.Next()
		} else {
			rl.mu.Unlock()
			c.JSON(http.StatusTooManyRequests, gin.H{
				"error": "Rate limit exceeded. Please try again later.",
			})
			c.Abort()
		}
	}
}

// Stop stops the cleanup ticker
func (rl *RateLimiter) Stop() {
	if rl.cleanupTicker != nil {
		rl.cleanupTicker.Stop()
	}
}