package middleware

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"sync"
	"time"
)

const (
	// loginFailureMaxKeys bounds how many distinct failure counters
	// each of the guard's two key sets holds. The submitted username
	// is part of a key, so the key set is attacker-influenced and
	// needs a hard cap or the limiter becomes the memory
	// amplification surface it exists to protect.
	//
	// A single-admin deployment has a handful of legitimate (client,
	// username) pairs, so 1024 is three orders of magnitude of
	// headroom before a real operator can be pushed onto the
	// fallback. It costs little: a counter is a ~64-byte key string,
	// a 32-byte window and map overhead, call it 170 bytes, so both
	// key sets full is 2 * 1024 * 170 bytes, under 0.4 MB.
	loginFailureMaxKeys = 1024

	// passwordVerifyConcurrency bounds how many Argon2id
	// verifications may run at once across every password-verifying
	// endpoint. Because credentials are now verified before any
	// limiter budget is spent, an attacker can force one hash per
	// request, and each hash allocates argon2Memory — 64 MB. Two
	// slots commit at most 128 MB to password hashing, which fits
	// inside the smallest container this service is realistically
	// given alongside its own working set; four would commit 256 MB
	// and crowd it. A single-admin product needs no concurrent
	// logins at all, so the second slot exists only so that one
	// stalled request does not serialise the endpoint.
	passwordVerifyConcurrency = 2

	// passwordVerifyWait is how long a request waits for a
	// verification slot before it is answered 503. Slots are handed
	// out in arrival order, so a legitimate request queues behind
	// the requests already waiting rather than behind the flood as a
	// whole. The wait is well inside the 60s request timeout.
	passwordVerifyWait = 5 * time.Second

	// passwordVerifyMaxWaiters bounds how many requests may be
	// queued for a slot at once. Past it, acquire sheds immediately
	// with 503 instead of joining the queue.
	//
	// The wait bounds how long one request occupies memory; this
	// bounds how many do so at the same time, and without it the
	// 128 MB hashing budget above is the smaller half of the real
	// footprint. At the 400 req/s a saturation attack can offer, an
	// unbounded queue would park ~2000 requests for the full five
	// seconds.
	//
	// A waiter costs far more than maxFormBodySize suggests: that
	// caps the raw body read, not what the parse retains. MaxBodySize,
	// CSRF and ParseForm all run before acquire, so a parked waiter
	// holds r.Form plus r.PostForm plus its header block for the
	// whole wait. Measured on the pinned go1.26.1 toolchain, as the
	// HeapAlloc delta across two GCs with 64 waiters parked in the
	// handler: an ordinary two-field login form retains ~0 MB, but a
	// 1 MB urlencoded body at Go's 10,000-parameter parse cap retains
	// 2.82 MB (3.09 MB with %41 escapes), and adding the ~0.9 MB of
	// headers httpMaxHeaderBytes allows takes it to 4.18 MB. The
	// retained parse and the header block dominate; the raw body does
	// not.
	//
	// Arithmetic, from the measured 4.18 MB worst case: 16 waiters
	// commit ~67 MB of queue memory, and peak commitment for the
	// endpoint is 128 MB of Argon2id plus the 18 requests that retain
	// a parsed form — 16 queued and the 2 being hashed — at
	// 18 * 4.18 MB, so ~75 MB: about 203 MB in all. Cross-check
	// against the deadline: two slots at the ~27 verifications/s
	// measured on a review host (with the race detector on, so the
	// real rate is higher) drain a full 16-deep queue in about 0.6 s,
	// far inside passwordVerifyWait.
	//
	// Those 203 MB are live bytes, not resident bytes: the Go
	// collector lets the heap reach roughly twice the live set before
	// collecting, with transient parse garbage on top. The review
	// measured a peak HeapAlloc of 392 MB against this guard under 18
	// adversarial requests, so provision on the order of 400 MB rather
	// than 203 MB.
	passwordVerifyMaxWaiters = 16

	// failureKeyHashBytes is how much of the username digest goes
	// into a failure key. 64 bits over at most loginFailureMaxKeys
	// live keys makes a collision negligible, and a collision would
	// only merge two usernames' failure counters, which throttles
	// sooner rather than later.
	failureKeyHashBytes = 8
)

// failureWindow counts failed credential verifications for one
// bucket, and records when that count lapses.
type failureWindow struct {
	count   int
	resetAt time.Time
}

// loginGuard is what replaced the pre-emptive rate limiter on the
// login POST.
//
// A limiter that spends budget on arrival cannot protect a
// single-admin product: behind the reverse proxy the deployment
// requires, with TRUSTED_PROXIES unset, every client keys on the
// proxy, so a stranger trickling five POSTs a minute keeps the one
// bucket full and the operator's own correct password is answered 429
// forever. There is no second administrative path.
//
// So budget is spent only by a FAILED verification. A correct
// password is never throttled, whatever the counters say, which is
// the only shape that guarantees the operator can get in. Two
// consequences follow and are handled here:
//
//   - Every login request now costs an Argon2id hash, so the number
//     running concurrently is bounded by slots. Without that bound
//     this trades an admin lockout for memory exhaustion, which is
//     strictly worse.
//   - Counting per (client, username) makes the key set
//     attacker-influenced, so both key sets are capped. Beyond the
//     per-username cap, failures fall back to a counter keyed on the
//     client alone; beyond that cap too, a failure is answered as
//     throttled without being recorded, since refusing to answer a
//     wrong password costs the operator nothing.
type loginGuard struct {
	mu     sync.Mutex
	byUser map[string]*failureWindow
	byAddr map[string]*failureWindow

	slots chan struct{}

	// queue holds one token per request waiting for a slot. A token
	// is taken non-blockingly, so a request that finds it full is
	// shed rather than queued, and is given up as soon as the wait
	// ends however it ends.
	queue chan struct{}

	limit    int
	interval time.Duration
	maxKeys  int
	wait     time.Duration

	// now is time.Now outside tests.
	now func() time.Time
}

// newLoginGuard builds a guard with the given failure limit per
// interval, key-set cap, verification concurrency, queue depth and
// slot wait.
func newLoginGuard(
	limit int,
	interval time.Duration,
	maxKeys, concurrency, maxWaiters int,
	wait time.Duration,
) *loginGuard {
	return &loginGuard{
		byUser:   make(map[string]*failureWindow),
		byAddr:   make(map[string]*failureWindow),
		slots:    make(chan struct{}, concurrency),
		queue:    make(chan struct{}, maxWaiters),
		limit:    limit,
		interval: interval,
		maxKeys:  maxKeys,
		wait:     wait,
		now:      time.Now,
	}
}

// acquire reserves a verification slot, waiting up to the guard's
// wait for one. It reports false when the queue of waiters is
// already full, when no slot became available in time, or when the
// request was cancelled first; the caller must then answer 503
// without verifying anything. The returned function releases the
// slot and must be called exactly once.
func (g *loginGuard) acquire(ctx context.Context) (func(), bool) {
	// A free slot is taken before any timer is armed, and before a
	// queue place is claimed: a request that never waits is not a
	// waiter. Without this preamble the bounded select below can find
	// its slot send and an already-expired timer ready at the same
	// time, and Go picks among ready cases uniformly at random — so a
	// process descheduled for longer than the wait sheds a request
	// with slots standing free, which is precisely when shedding is
	// least defensible.
	//
	// This cannot let a late arrival barge past a queued waiter. A
	// waiter can only be parked on a FULL buffer, and a release
	// refills that buffer from the head of the send queue under the
	// channel lock, so the buffer never appears non-full while anyone
	// is parked and this send fails whenever there is a waiter.
	select {
	case g.slots <- struct{}{}:
		return func() { <-g.slots }, true
	default:
	}

	// Shedding past the queue depth is what keeps waiting memory
	// bounded; the wait alone only bounds how long one waiter holds
	// its parsed form, not how many hold one at once.
	select {
	case g.queue <- struct{}{}:
	default:
		return nil, false
	}

	// Held only for the wait. A request that gets a slot gives its
	// queue token back before it starts hashing, so the depth is a
	// bound on waiters rather than on requests in the handler.
	defer func() { <-g.queue }()

	timer := time.NewTimer(g.wait)
	defer timer.Stop()

	// The blocking send is deliberate: a receive on a full buffered
	// channel hands the slot straight to the head of the send queue,
	// so slots go out in arrival order and a later arrival cannot
	// barge past a request already waiting.
	select {
	case g.slots <- struct{}{}:
		return func() { <-g.slots }, true
	case <-timer.C:
		return nil, false
	case <-ctx.Done():
		return nil, false
	}
}

// fail records one failed credential verification by clientKey
// against username, and reports whether this client has now spent
// its failure budget and should be answered 429.
func (g *loginGuard) fail(clientKey, username string) bool {
	g.mu.Lock()
	defer g.mu.Unlock()

	now := g.now()

	window := g.window(
		g.byUser, userFailureKey(clientKey, username), now,
	)
	if window == nil {
		window = g.window(g.byAddr, clientKey, now)
	}

	if window == nil {
		// Both key sets are full and neither already tracks this
		// client, so nothing can be counted without unbounded
		// growth. Answering the failure as throttled is the safe
		// direction: it never touches a correct password.
		return true
	}

	window.count++

	return window.count >= g.limit
}

// succeed forgives clientKey's failures against username. A correct
// password clears the counters, so an operator who mistypes several
// times and then gets it right is not throttled afterwards.
func (g *loginGuard) succeed(clientKey, username string) {
	g.mu.Lock()
	defer g.mu.Unlock()

	delete(g.byUser, userFailureKey(clientKey, username))
	delete(g.byAddr, clientKey)
}

// window returns the live counter for key in set, resetting a lapsed
// one and creating a missing one when the cap allows. It returns nil
// only when key is absent and set is full even after lapsed entries
// are swept.
func (g *loginGuard) window(
	set map[string]*failureWindow,
	key string,
	now time.Time,
) *failureWindow {
	window, ok := set[key]
	if ok {
		if !now.Before(window.resetAt) {
			window.count = 0
			window.resetAt = now.Add(g.interval)
		}

		return window
	}

	if len(set) >= g.maxKeys {
		sweepLapsed(set, now)
	}

	if len(set) >= g.maxKeys {
		return nil
	}

	window = &failureWindow{resetAt: now.Add(g.interval)}
	set[key] = window

	return window
}

// sweepLapsed drops counters whose interval has elapsed.
func sweepLapsed(set map[string]*failureWindow, now time.Time) {
	for key, window := range set {
		if !now.Before(window.resetAt) {
			delete(set, key)
		}
	}
}

// userFailureKey identifies one (client, submitted username) pair.
// The username is hashed rather than embedded: a submitted username
// is attacker-controlled text of attacker-chosen length, and hashing
// makes every key the same size whatever was sent.
func userFailureKey(clientKey, username string) string {
	sum := sha256.Sum256([]byte(username))

	return clientKey + "|" +
		hex.EncodeToString(sum[:failureKeyHashBytes])
}

// guard returns the middleware's login guard, building it on first
// use so that every construction path — fx and the test constructor
// alike — gets one.
func (m *Middleware) guard() *loginGuard {
	m.loginGuardOnce.Do(func() {
		m.loginGuard = newLoginGuard(
			loginRateLimit,
			loginRateInterval,
			loginFailureMaxKeys,
			passwordVerifyConcurrency,
			passwordVerifyMaxWaiters,
			passwordVerifyWait,
		)
	})

	return m.loginGuard
}

// BeginPasswordVerification reserves one of the bounded Argon2id
// verification slots. It reports false when the queue of waiting
// requests is already at passwordVerifyMaxWaiters, or when no slot
// became free within passwordVerifyWait; in either case the caller
// must answer 503 and must not verify a password. The returned
// function releases the slot and must be called exactly once.
//
// Every endpoint that hashes a password on request must go through
// this, or the bound has a hole: the memory is committed per hash,
// not per endpoint.
func (m *Middleware) BeginPasswordVerification(
	ctx context.Context,
) (func(), bool) {
	return m.guard().acquire(ctx)
}

// RecordLoginFailure counts a failed credential verification for the
// request's client against the submitted username, and reports
// whether the response should be 429 rather than 401.
func (m *Middleware) RecordLoginFailure(
	r *http.Request,
	username string,
) bool {
	throttled := m.guard().fail(m.clientKey(r), username)
	if throttled {
		m.log.Warn(
			"login failure limit exceeded", "path", r.URL.Path,
		)
	}

	return throttled
}

// ForgiveLoginFailures clears the failure counters for the request's
// client and the submitted username after a successful
// authentication.
func (m *Middleware) ForgiveLoginFailures(
	r *http.Request,
	username string,
) {
	m.guard().succeed(m.clientKey(r), username)
}

// LoginFailureInterval is how long a spent login failure budget
// takes to refill, which is what a throttled login answers as
// Retry-After.
func (m *Middleware) LoginFailureInterval() time.Duration {
	return m.guard().interval
}