package middleware

import (
	"sync"
	"time"
)

// FailureLimiter throttles repeated authentication failures per key (usually a
// client IP).
//
// Deliberately not a request limiter: a WebDAV sync client legitimately sends
// hundreds of requests per minute, so capping requests would break normal use.
// Only failed logins count, and a single success clears the record — which
// leaves honest clients untouched while a password guesser runs out of tries.
type FailureLimiter struct {
	mu      sync.Mutex
	entries map[string]*failureEntry

	maxFailures int
	window      time.Duration
	lockout     time.Duration
	lastSweep   time.Time
}

type failureEntry struct {
	count       int
	windowStart time.Time
	lockedUntil time.Time
}

// NewFailureLimiter allows maxFailures within window; exceeding that locks the
// key out for lockout.
func NewFailureLimiter(maxFailures int, window, lockout time.Duration) *FailureLimiter {
	return &FailureLimiter{
		entries:     make(map[string]*failureEntry),
		maxFailures: maxFailures,
		window:      window,
		lockout:     lockout,
		lastSweep:   time.Now(),
	}
}

// Allowed reports whether the key may attempt authentication. The second return
// value is how long the caller must wait, for a Retry-After header.
func (l *FailureLimiter) Allowed(key string) (bool, time.Duration) {
	l.mu.Lock()
	defer l.mu.Unlock()

	e, ok := l.entries[key]
	if !ok {
		return true, 0
	}
	now := time.Now()
	if now.Before(e.lockedUntil) {
		return false, e.lockedUntil.Sub(now)
	}
	return true, 0
}

// Fail records a failed attempt and locks the key out once the limit is hit.
func (l *FailureLimiter) Fail(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	l.sweepLocked(now)

	e, ok := l.entries[key]
	if !ok || now.Sub(e.windowStart) > l.window {
		l.entries[key] = &failureEntry{count: 1, windowStart: now}
		return
	}
	e.count++
	if e.count >= l.maxFailures {
		e.lockedUntil = now.Add(l.lockout)
		// Restart the window so the lockout does not extend forever on the
		// retries that arrive while it is already active.
		e.count = 0
		e.windowStart = now
	}
}

// Reset clears the record after a successful authentication.
func (l *FailureLimiter) Reset(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.entries, key)
}

// sweepLocked drops stale records so the map cannot grow without bound from
// one-off failures across many addresses. Caller must hold the mutex.
func (l *FailureLimiter) sweepLocked(now time.Time) {
	if now.Sub(l.lastSweep) < l.window {
		return
	}
	l.lastSweep = now
	for k, e := range l.entries {
		if now.After(e.lockedUntil) && now.Sub(e.windowStart) > l.window {
			delete(l.entries, k)
		}
	}
}