package wsauth

import (
	"sync"
	"time"
)

// RateLimiter is a sliding-window per-key counter used to throttle login
// attempts. Two instances are typically created: one keyed by client IP
// (to slow distributed brute force), one keyed by username (to slow
// targeted attacks against a known account).
//
// Design notes:
//   - Denied attempts are NOT recorded — the window slides naturally and a
//     legitimate user who fat-fingers their password recovers as soon as
//     the oldest attempt ages out, while a determined attacker is capped
//     at `limit` successful attempts per `window` indefinitely.
//   - Lazy cleanup: stale timestamps for a key are pruned on every Allow()
//     call. Truly idle keys are GC'd by Sweep(), which callers should run
//     periodically from a background goroutine.
//   - Map size is bounded by the count of recently-active keys; for the
//     web UI's expected load (a handful of users + occasional scanners),
//     no extra GC pressure considerations needed.
type RateLimiter struct {
	mu      sync.Mutex
	limit   int
	window  time.Duration
	entries map[string][]time.Time
}

// NewRateLimiter returns a limiter that allows up to `limit` events per
// `window` duration per key. Zero or negative limit/window disables the
// limiter (Allow always returns true) — useful for tests / dev mode.
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	return &RateLimiter{
		limit:   limit,
		window:  window,
		entries: make(map[string][]time.Time),
	}
}

// Allow records an attempt for `key` if and only if the caller is under
// the per-key limit. Returns true when allowed, false when over limit.
// Empty key is treated as "no throttle" (returns true without recording)
// so the caller can fall through when the IP/username is unavailable.
func (r *RateLimiter) Allow(key string) bool {
	if r == nil || r.limit <= 0 || r.window <= 0 || key == "" {
		return true
	}
	r.mu.Lock()
	defer r.mu.Unlock()

	cutoff := time.Now().Add(-r.window)
	times := r.entries[key]
	// Compact in place — keep only timestamps within the window.
	keep := times[:0]
	for _, t := range times {
		if t.After(cutoff) {
			keep = append(keep, t)
		}
	}

	if len(keep) >= r.limit {
		// Update the map even when denying so the compacted slice doesn't
		// keep stale entries forever. Don't append the new attempt: that
		// would let attackers extend the window arbitrarily.
		r.entries[key] = keep
		return false
	}

	r.entries[key] = append(keep, time.Now())
	return true
}

// Reset clears state for a key. Call on successful login to give the user
// a fresh budget — otherwise a string of failed attempts followed by a
// correct one still leaves the budget partially consumed.
func (r *RateLimiter) Reset(key string) {
	if r == nil || key == "" {
		return
	}
	r.mu.Lock()
	delete(r.entries, key)
	r.mu.Unlock()
}

// Sweep removes entries whose timestamps have all aged out of the window.
// Safe to call concurrently with Allow. Intended for periodic invocation
// from a background ticker (e.g. every window-length) to bound the map.
func (r *RateLimiter) Sweep() {
	if r == nil {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	cutoff := time.Now().Add(-r.window)
	for key, times := range r.entries {
		keep := times[:0]
		for _, t := range times {
			if t.After(cutoff) {
				keep = append(keep, t)
			}
		}
		if len(keep) == 0 {
			delete(r.entries, key)
		} else {
			r.entries[key] = keep
		}
	}
}