package idear

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

// RateLimit bounds how often ONE client may hit idear's two public
// routes. Both of them read a secret: GET /invitations/{token} answers
// questions about an invitation to anyone holding the token, and POST
// /invitations/{token} spends one. An unauthenticated lookup of a
// secret that answers as fast as the network allows is a free oracle,
// so the limiter is not optional and there is no way to switch it off
// — only to widen it.
//
// The defaults are a token bucket of Burst requests that refills one
// token every Every: twenty at once, then one every three seconds
// (twenty a minute sustained). A person opening their invitation link
// and posting it never notices; a script walking the 2^256 token space
// gets twenty guesses a minute out of one address, which is not a
// meaningful improvement on none.
type RateLimit struct {
	// Burst is how many requests a client may make back to back.
	// Default 20.
	Burst int

	// Every is how long one token takes to come back. Default 3s.
	Every time.Duration

	// Max is the ceiling on how many clients are tracked at once —
	// the memory bound. Default 4096.
	//
	// When the table is full, idear first drops every bucket that has
	// refilled to Burst (such a bucket is indistinguishable from a
	// client that has never been seen, so dropping it costs nothing
	// and forgives nobody). If it is STILL full, the request is
	// REFUSED rather than admitted: on a route whose whole job is to
	// slow down guessing, an overflowing table is evidence of the
	// attack the limiter exists for, and failing open there would
	// make Max the way around the limit. The cost is that a wide
	// enough spray can lock out real invitees for as long as it lasts
	// — stated so it is a decision and not a surprise.
	//
	// A routed IPv6 allocation (a /48 or /56) can supply thousands of
	// the distinct keys this table tracks from ONE attacker; see
	// clientIP for the sizing this implies for Max.
	Max int
}

const (
	defaultRateBurst = 20
	defaultRateEvery = 3 * time.Second
	defaultRateMax   = 4096
)

// limiter is the bounded in-memory token bucket behind RateLimit. One
// per *Handlers, so it is per-process and NOT shared between replicas:
// two instances behind a load balancer give a client two budgets. That
// is a real weakening and it is accepted, because the alternative is a
// database write on every unauthenticated request — which is the
// resource the limiter is trying to protect.
type limiter struct {
	burst float64
	every time.Duration
	max   int

	// now is the clock, injectable so a test can prove the refill
	// without sleeping through it.
	now func() time.Time

	mu      sync.Mutex
	buckets map[string]*bucket
}

type bucket struct {
	tokens float64
	seen   time.Time
}

func newLimiter(rl RateLimit) *limiter {
	if rl.Burst <= 0 {
		rl.Burst = defaultRateBurst
	}
	if rl.Every <= 0 {
		rl.Every = defaultRateEvery
	}
	if rl.Max <= 0 {
		rl.Max = defaultRateMax
	}
	return &limiter{
		burst:   float64(rl.Burst),
		every:   rl.Every,
		max:     rl.Max,
		now:     func() time.Time { return time.Now() },
		buckets: make(map[string]*bucket),
	}
}

// allow spends one token for key, reporting whether the request may
// proceed. It is safe for concurrent use: the two public routes are
// exactly where concurrent unauthenticated traffic arrives.
func (l *limiter) allow(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := l.now()
	b, ok := l.buckets[key]
	if !ok {
		if len(l.buckets) >= l.max {
			l.sweep(now)
		}
		if len(l.buckets) >= l.max {
			// Fail closed. See RateLimit.Max.
			return false
		}
		b = &bucket{tokens: l.burst, seen: now}
		l.buckets[key] = b
	} else {
		l.refill(b, now)
	}
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

// refill credits a bucket for the time since it was last touched,
// capped at burst.
func (l *limiter) refill(b *bucket, now time.Time) {
	if elapsed := now.Sub(b.seen); elapsed > 0 {
		b.tokens += float64(elapsed) / float64(l.every)
		if b.tokens > l.burst {
			b.tokens = l.burst
		}
	}
	b.seen = now
}

// sweep drops every bucket that has refilled to full. A full bucket
// carries no state a fresh one would not have, so this frees memory
// without forgiving anybody: a client mid-penalty is never swept.
//
// Called only when the table is at Max, which is why there is no
// background goroutine here — a limiter with no traffic needs no
// sweeper, and a *Handlers must not leak one for the life of the
// process.
func (l *limiter) sweep(now time.Time) {
	for k, b := range l.buckets {
		l.refill(b, now)
		if b.tokens >= l.burst {
			delete(l.buckets, k)
		}
	}
}

// size is the number of tracked clients — for the test that proves the
// table is actually bounded.
func (l *limiter) size() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.buckets)
}

// clientIP is the default rate-limit key: the client's NETWORK, not
// its address.
//
// It is "per-IP-ish" and not per-person on purpose — there is nobody
// to identify on an unauthenticated route.
//
// THE FOLD TO /64 IS THE WHOLE POINT, and keying on the bare address
// was a defect rather than a simplification. Every IPv6 host is handed
// a /64: that is 2^64 source addresses, free to rotate, one per
// request. Keyed on the address, each one is an unseen client with a
// fresh burst, so the per-client limit does not exist at all — and it
// is worse than a limit that merely leaks. Folding to the /64 gives a
// SINGLE-/64 attacker one bucket, which is what it should always have
// had.
//
// It does not give a ROUTED attacker one bucket. A /48 or /56 —
// ordinary residential and business allocations, not an exotic
// stretch — carries 65,536 or 16,384 distinct /64s respectively, and
// every one folds to a different key. That is easily enough to fill
// a default-sized (4,096-entry) table on its own, at which point
// allow FAILS CLOSED and refuses every UNSEEN client on the table:
// real invitees, and orphans, whose only healing path is
// POST /invitations/{token}. One machine holding a /56 can hold an
// instance in that state indefinitely. Failing closed is still the
// right direction on a guessing-slowdown; the /64 fold is still
// strictly better than keying on the bare address (2^64 keys instead
// of 2^16). It is not, by itself, enough to assume one attacker means
// one bucket.
//
// Size RateLimit.Max with that in mind: it bounds MEMORY, not
// attacker-controlled prefixes, so raising it only buys headroom
// against a /48-or-wider spray, never immunity from one — the fold is
// fixed at /64 and is not configurable. An app that expects hostile
// IPv6 traffic should set Max well past the default (each entry is a
// handful of words, so a table sized in the hundreds of thousands
// costs low-single-digit megabytes) and treat the fail-closed path as
// a real operational outcome to monitor for, not a theoretical one.
//
// IPv4 is used whole (a /32), which is the same rule read the same
// way: the smallest unit an operator is routinely handed. Folding
// IPv4 further would put a whole CGNAT or campus behind one bucket.
// IPv4-mapped v6 addresses (::ffff:203.0.113.9) are unmapped first,
// so one client cannot hold two budgets by switching representation.
//
// Behind a reverse proxy EVERY request arrives from the proxy's
// address and they would share a single bucket, which turns the
// limiter into a global one and locks out real invitees; an app in
// that shape must set HandlerConfig.ClientKey to read its own TRUSTED
// forwarding header. idear never reads X-Forwarded-For (or any other
// forwarding header) itself, because a header idear cannot verify is
// a header an attacker can spoof to mint unlimited budgets — the
// exact failure the /64 fold exists to prevent, handed over for free.
// TestClientIPIgnoresForwardingHeaders pins that.
func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		// No port: a unix socket, or a test. Use it whole rather than
		// dropping it, so it still keys to something.
		host = r.RemoteAddr
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return host
	}
	// The zone ("fe80::1%eth0") is the local interface, not the
	// client, and would split one client's budget in two.
	addr = addr.WithZone("")
	if addr.Is4() || addr.Is4In6() {
		return addr.Unmap().String()
	}
	prefix, err := addr.Prefix(64)
	if err != nil {
		return addr.String()
	}
	return prefix.String()
}