package auth

import (
	"sync"
	"time"

	"github.com/iceBear67/simplepages/internal/cache"
)

// Limiter throttles clients that keep presenting bad credentials.
//
// It exists for CPU, not for guessing. An 80-bit key id plus a 256-bit secret
// cannot be brute forced, so the realistic attack is not "eventually guess a
// token" but "make the server parse, look up and hash forever". Only failed
// attempts consume budget, so a busy CI runner pushing hundreds of valid
// requests a second is never touched.
//
// The bucket map is an LRU with a hard cap, because its keys are whatever
// addresses show up: an unbounded map keyed by attacker-chosen input would turn
// the defence into a memory exhaustion vector of its own.
type Limiter struct {
	buckets *cache.Cache[string, *bucket]
	burst   float64
	refill  float64 // tokens per second

	now func() time.Time // swapped in tests
}

type bucket struct {
	mu     sync.Mutex
	tokens float64
	last   time.Time
}

// NewLimiter allows burst consecutive failures per client, refilling to full
// over period, and tracks at most maxClients addresses.
func NewLimiter(burst int, period time.Duration, maxClients int) *Limiter {
	if burst < 1 {
		burst = 1
	}
	if period <= 0 {
		period = time.Minute
	}
	// Idle buckets are dropped after twice the refill period: by then a bucket
	// has refilled completely, so forgetting it and recreating it full are the
	// same thing.
	return &Limiter{
		buckets: cache.New[string, *bucket](maxClients, 2*period),
		burst:   float64(burst),
		refill:  float64(burst) / period.Seconds(),
		now:     time.Now,
	}
}

// Allow reports whether client has any budget left. It does not consume any:
// a request that turns out to authenticate correctly should cost nothing.
func (l *Limiter) Allow(client string) bool {
	if l == nil {
		return true
	}
	b := l.bucket(client)
	b.mu.Lock()
	defer b.mu.Unlock()
	l.refillLocked(b)
	return b.tokens >= 1
}

// Fail records one failed attempt for client.
func (l *Limiter) Fail(client string) {
	if l == nil {
		return
	}
	b := l.bucket(client)
	b.mu.Lock()
	defer b.mu.Unlock()
	l.refillLocked(b)
	if b.tokens >= 1 {
		b.tokens--
	} else {
		b.tokens = 0
	}
}

// RetryAfter estimates how long client must wait for one token, for the
// Retry-After header. It rounds up to whole seconds, and never returns zero
// while the client is actually throttled.
func (l *Limiter) RetryAfter(client string) time.Duration {
	if l == nil {
		return 0
	}
	b := l.bucket(client)
	b.mu.Lock()
	defer b.mu.Unlock()
	l.refillLocked(b)
	if b.tokens >= 1 {
		return 0
	}
	need := 1 - b.tokens
	d := time.Duration(need / l.refill * float64(time.Second))
	return d.Round(time.Second) + time.Second
}

func (l *Limiter) bucket(client string) *bucket {
	return l.buckets.GetOrCreate(client, func() *bucket {
		return &bucket{tokens: l.burst, last: l.now()}
	})
}

func (l *Limiter) refillLocked(b *bucket) {
	now := l.now()
	if elapsed := now.Sub(b.last); elapsed > 0 {
		b.tokens += elapsed.Seconds() * l.refill
		if b.tokens > l.burst {
			b.tokens = l.burst
		}
		b.last = now
	}
}