package mining

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

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

// rl := NewRateLimiter(10, 20) // 10 req/s, burst 20
// router.Use(rl.Middleware())
type RateLimiter struct {
	requestsPerSecond int
	burst             int
	clients           map[string]*rateLimitClient
	mutex             sync.RWMutex
	stopChan          chan struct{}
	stopped           bool
}

type rateLimitClient struct {
	tokens    float64
	lastCheck time.Time
}

// rl := NewRateLimiter(10, 20)  // 10 requests/second, burst of 20
// defer rl.Stop()
func NewRateLimiter(requestsPerSecond, burst int) *RateLimiter {
	rl := &RateLimiter{
		requestsPerSecond: requestsPerSecond,
		burst:             burst,
		clients:           make(map[string]*rateLimitClient),
		stopChan:          make(chan struct{}),
	}

	// Start cleanup goroutine
	go rl.cleanupLoop()

	return rl
}

// go rl.cleanupLoop()  // started by NewRateLimiter; runs until rl.Stop()
func (rl *RateLimiter) cleanupLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-rl.stopChan:
			return
		case <-ticker.C:
			rl.cleanup()
		}
	}
}

// rl.cleanup()  // called every minute by cleanupLoop; evicts IPs idle for >5 minutes
func (rl *RateLimiter) cleanup() {
	rl.mutex.Lock()
	defer rl.mutex.Unlock()

	for ip, client := range rl.clients {
		if time.Since(client.lastCheck) > 5*time.Minute {
			delete(rl.clients, ip)
		}
	}
}

// rl.Stop()  // call on shutdown to release the cleanup goroutine
func (rl *RateLimiter) Stop() {
	rl.mutex.Lock()
	defer rl.mutex.Unlock()

	if !rl.stopped {
		close(rl.stopChan)
		rl.stopped = true
	}
}

// router.Use(rl.Middleware())  // install before route handlers
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()

		rl.mutex.Lock()
		client, exists := rl.clients[ip]
		if !exists {
			client = &rateLimitClient{tokens: float64(rl.burst), lastCheck: time.Now()}
			rl.clients[ip] = client
		}

		// Token bucket algorithm
		now := time.Now()
		elapsed := now.Sub(client.lastCheck).Seconds()
		client.tokens += elapsed * float64(rl.requestsPerSecond)
		if client.tokens > float64(rl.burst) {
			client.tokens = float64(rl.burst)
		}
		client.lastCheck = now

		if client.tokens < 1 {
			rl.mutex.Unlock()
			respondWithError(c, http.StatusTooManyRequests, "RATE_LIMITED",
				"too many requests", "rate limit exceeded")
			c.Abort()
			return
		}

		client.tokens--
		rl.mutex.Unlock()
		c.Next()
	}
}

// if rl.ClientCount() == 0 { /* no active clients */ }
func (rl *RateLimiter) ClientCount() int {
	rl.mutex.RLock()
	defer rl.mutex.RUnlock()
	return len(rl.clients)
}