// Package ratelimit is a dependency free, in memory token bucket rate limiter keyed by client IP. It
// protects the auth endpoints from credential stuffing and the API from a single abusive caller. It is
// a defense in depth layer, not a substitute for a WAF or a per tenant quota at the edge; production is
// expected to also sit behind a reverse proxy that rate limits. Idle buckets are swept so memory stays
// bounded under a changing IP set.
package ratelimit

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

type bucket struct {
	tokens float64
	last   time.Time
}

// Limiter is a concurrency safe token bucket keyed by an arbitrary string.
type Limiter struct {
	mu      sync.Mutex
	buckets map[string]*bucket
	rate    float64 // tokens refilled per second
	burst   float64 // bucket capacity
	now     func() time.Time
}

// New builds a limiter that refills ratePerSec tokens per second up to burst capacity.
func New(ratePerSec float64, burst int) *Limiter {
	return &Limiter{buckets: make(map[string]*bucket), rate: ratePerSec, burst: float64(burst), now: time.Now}
}

// Allow consumes one token for key, returning false when the bucket is empty (request should be denied).
func (l *Limiter) Allow(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()
	b := l.buckets[key]
	if b == nil {
		b = &bucket{tokens: l.burst, last: now}
		l.buckets[key] = b
	}
	if elapsed := now.Sub(b.last).Seconds(); elapsed > 0 {
		b.tokens += elapsed * l.rate
		if b.tokens > l.burst {
			b.tokens = l.burst
		}
		b.last = now
	}
	if b.tokens >= 1 {
		b.tokens--
		return true
	}
	return false
}

// sweep drops buckets idle longer than ttl that have fully recovered, so dropping them loses no state
// (a fresh bucket would start full anyway). A still throttled bucket (not yet refilled to capacity) is
// kept so its window is enforced. Refill is lazy, so it computes the would be current token count
// rather than reading the stale stored value.
func (l *Limiter) sweep(ttl time.Duration) {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()
	cutoff := now.Add(-ttl)
	for k, b := range l.buckets {
		if !b.last.Before(cutoff) {
			continue // recently active, keep
		}
		if b.tokens+now.Sub(b.last).Seconds()*l.rate >= l.burst {
			delete(l.buckets, k)
		}
	}
}

// StartSweeper runs sweep every interval until stop is closed. Call once from main.
func (l *Limiter) StartSweeper(stop <-chan struct{}, interval, ttl time.Duration) {
	go func() {
		t := time.NewTicker(interval)
		defer t.Stop()
		for {
			select {
			case <-stop:
				return
			case <-t.C:
				l.sweep(ttl)
			}
		}
	}()
}

// Middleware limits by client IP. trustedHeader, when non empty, is the header a trusted reverse proxy
// sets with the real client IP (for example X-Real-IP); otherwise the TCP peer is used. On the limit it
// returns 429 with a short Retry-After and does not call the next handler.
func (l *Limiter) Middleware(trustedHeader string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if !l.Allow(ClientIP(r, trustedHeader)) {
				w.Header().Set("Retry-After", "1")
				http.Error(w, "rate limited", http.StatusTooManyRequests)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// ClientIP resolves the caller IP. It trusts trustedHeader ONLY when set (an operator opts in because
// the app sits behind a proxy that overwrites it); otherwise it uses the TCP peer, which cannot be
// spoofed. A comma list (proxy chain) takes the first entry, the original client.
func ClientIP(r *http.Request, trustedHeader string) string {
	if trustedHeader != "" {
		if v := strings.TrimSpace(r.Header.Get(trustedHeader)); v != "" {
			if i := strings.IndexByte(v, ','); i >= 0 {
				v = strings.TrimSpace(v[:i])
			}
			return v
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}