package auth

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

// RateLimiter is a per-key token bucket. Keys are typically "<route>:<ip>"
// composed by the caller. The limiter is in-memory; on a restart the buckets
// reset, which is acceptable at single-operator scale.
type RateLimiter struct {
	mu      sync.Mutex
	buckets map[string]*bucket
	limit   int           // tokens per window
	window  time.Duration // refill period
}

type bucket struct {
	tokens     int
	lastRefill time.Time
}

// NewRateLimiter constructs a limiter that grants `limit` tokens per `window`
// per key. Defaults are 5/5min — calling code overrides for the stricter
// recovery endpoint.
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	if limit <= 0 {
		limit = 5
	}
	if window <= 0 {
		window = 5 * time.Minute
	}
	return &RateLimiter{
		buckets: map[string]*bucket{},
		limit:   limit,
		window:  window,
	}
}

// Allow returns true if the call should proceed. It atomically decrements the
// key's bucket; an empty bucket returns false.
func (r *RateLimiter) Allow(key string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()

	now := time.Now()
	b, ok := r.buckets[key]
	if !ok {
		b = &bucket{tokens: r.limit, lastRefill: now}
		r.buckets[key] = b
	}
	// Full refill once per window — simpler than per-token drip and gives
	// the operator a clean "wait until X" UX when rate-limited.
	if now.Sub(b.lastRefill) >= r.window {
		b.tokens = r.limit
		b.lastRefill = now
	}
	if b.tokens <= 0 {
		return false
	}
	b.tokens--
	return true
}

// KeyForRoute composes a route-scoped rate-limit key from the request's IP.
// X-Forwarded-For / X-Real-IP are handled by chi's RealIP middleware before
// we see the request.
func KeyForRoute(route string, r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	return route + ":" + host
}