package auth

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

	"github.com/danielgtaylor/huma/v2"
)

// RateLimiter provides per-IP rate limiting for authenticated API endpoints,
// complementing the login-specific failLimiter with a broader cap on all
// requests. The window aligns to wall-clock intervals so all IPs share the
// same boundary.
type RateLimiter struct {
	mu       sync.Mutex
	buckets  map[string]*tokenBucket
	limit    int
	interval time.Duration
}

type tokenBucket struct {
	count     int
	windowEnd time.Time
}

const maxRateLimitKeys = 10000

func NewRateLimiter(limit int, interval time.Duration) *RateLimiter {
	return &RateLimiter{
		buckets:  map[string]*tokenBucket{},
		limit:    limit,
		interval: interval,
	}
}

// Allow reports whether ip may make a request now. Returns false when the
// limit is exceeded OR the map is full (fail-closed).
func (l *RateLimiter) Allow(ip string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	b, ok := l.buckets[ip]
	if !ok || now.After(b.windowEnd) {
		if len(l.buckets) >= maxRateLimitKeys {
			return false
		}
		l.buckets[ip] = &tokenBucket{count: 1, windowEnd: now.Add(l.interval)}
		return true
	}
	if b.count >= l.limit {
		return false
	}
	b.count++
	return true
}

// RateLimitMiddleware returns Huma middleware that rejects requests exceeding
// the per-IP rate limit with 429 Too Many Requests. It runs before the RBAC
// check so abusive IPs are dropped early. The IP is read from the context set
// by WithClientIP; if absent the request passes through unthrottled.
func RateLimitMiddleware(api huma.API, rl *RateLimiter) func(huma.Context, func(huma.Context)) {
	return func(ctx huma.Context, next func(huma.Context)) {
		ip := ClientIP(ctx.Context())
		if ip != "" && !rl.Allow(ip) {
			huma.WriteErr(api, ctx, http.StatusTooManyRequests, "rate limit exceeded, try again later")
			return
		}
		next(ctx)
	}
}
