package auth

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

// failLimiter throttles repeated login failures keyed by "username|ip". After
// max failures it imposes a cooldown of window before any further attempt for
// that key is accepted, blunting brute force against PAM/shadow.
//
// ponytail: in-memory, single-process, fixed cooldown - correct for a single
// app instance. State is lost on restart, but only an operator restarts the
// process (an attacker can't), so persisting it would buy nothing. Source
// spoofing is handled by the network (VPN-only + trusted proxy sets XFF), not
// here. Reach for pam_faillock only if a failed web login should also lock the
// OS account against ssh/console - a different layer we deliberately don't span.
type failLimiter struct {
	mu       sync.Mutex
	attempts map[string]*attemptState
	max      int
	window   time.Duration
}

type attemptState struct {
	count int
	until time.Time
}

// maxTrackedKeys bounds memory: an attacker rotating username/IP can't grow the
// map without limit. When exceeded we drop all throttle state - a crude reset
// that briefly forgets cooldowns, acceptable for a single-node panel.
const maxTrackedKeys = 10000

func newFailLimiter(max int, window time.Duration) *failLimiter {
	return &failLimiter{attempts: map[string]*attemptState{}, max: max, window: window}
}

// blocked reports whether the key is currently in cooldown.
func (l *failLimiter) blocked(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	s := l.attempts[key]
	return s != nil && time.Now().Before(s.until)
}

// fail records a failed attempt and starts a cooldown once max is reached.
func (l *failLimiter) fail(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if len(l.attempts) > maxTrackedKeys {
		l.attempts = map[string]*attemptState{}
	}
	s := l.attempts[key]
	if s == nil {
		s = &attemptState{}
		l.attempts[key] = s
	}
	s.count++
	if s.count >= l.max {
		s.until = time.Now().Add(l.window)
		s.count = 0 // restart the window after the cooldown is set
	}
}

// reset clears state for a key after a successful login.
func (l *failLimiter) reset(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.attempts, key)
}

type ctxKey int

const clientIPKey ctxKey = 0

// WithClientIP wraps a handler so the client's IP is available downstream via
// ClientIP, so the login throttle can key on source IP without coupling to a
// specific HTTP adapter.
//
// When trustProxy is set, the IP is taken from the first X-Forwarded-For hop
// (the original client behind the reverse proxy). X-Forwarded-For is fully
// caller-controlled, so this is only honored when an admin has opted in by
// putting nadir behind a proxy - and nadir must then be reachable only by that
// proxy. When trustProxy is false the header is ignored and RemoteAddr wins.
func WithClientIP(trustProxy bool, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ip, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			ip = r.RemoteAddr
		}
		if trustProxy {
			if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
				first, _, _ := strings.Cut(xff, ",")
				ip = strings.TrimSpace(first)
			}
		}
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), clientIPKey, ip)))
	})
}

// ClientIP returns the IP recorded by WithClientIP, or "" if absent.
func ClientIP(ctx context.Context) string {
	ip, _ := ctx.Value(clientIPKey).(string)
	return ip
}