package bot

import (
	"sync"
	"time"
)

// The public bot answers anyone, so one person with a script can spend the
// whole service's budget against the cep.is API. These constants bound what a
// single Telegram account may cost.
//
// Burst is what a human actually does: paste a handful of CEPs in a row, then
// think. Refill is the sustained rate afterwards — 20 lookups a minute, which
// no person reaches by hand and which a loop hits within a second.
//
// WarnInterval keeps the refusal from becoming its own flood: someone held at
// the limit is told once, then ignored until the window passes, so a script
// hammering the bot does not turn into the bot hammering Telegram.
const (
	Burst        = 5
	Refill       = 3 * time.Second
	WarnInterval = 30 * time.Second
)

// The bucket table is swept rather than kept forever: every distinct user gets
// an entry, and a public bot meets a lot of users. An entry is dropped once it
// has been idle long enough that it would have refilled to full anyway, at
// which point forgetting it and recreating it are the same thing.
const (
	bucketIdleTTL = 15 * time.Minute
	sweepEvery    = 5 * time.Minute
)

// Decision is the answer of [Limiter.Allow].
type Decision struct {
	// Allowed reports whether the request may proceed.
	Allowed bool

	// Warn reports whether the user should be told they are being limited.
	// It is true at most once per [WarnInterval] per key.
	Warn bool

	// RetryAfter is roughly how long until the next token, meant for the
	// message shown to the user.
	RetryAfter time.Duration
}

type bucket struct {
	tokens   float64
	lastSeen time.Time
	lastWarn time.Time
}

// Limiter is a token bucket per key, where the key is the Telegram user id — or
// the chat id for channel posts, which carry no author.
//
// It is in-memory and per-process on purpose: the bot runs as a single
// container, and a shared store would add a dependency and a failure mode to
// protect against something a single process can already see.
type Limiter struct {
	burst        float64
	refill       time.Duration
	warnInterval time.Duration
	now          func() time.Time

	mu        sync.Mutex
	buckets   map[int64]*bucket
	lastSweep time.Time
}

// NewLimiter builds a limiter with the package constants.
func NewLimiter() *Limiter {
	return NewLimiterWith(Burst, Refill, WarnInterval, time.Now)
}

// NewLimiterWith builds a limiter with explicit parameters and an injectable
// clock, which is how the test suite exercises refilling without waiting for it.
func NewLimiterWith(burst int, refill, warnInterval time.Duration, now func() time.Time) *Limiter {
	if now == nil {
		now = time.Now
	}
	return &Limiter{
		burst:        float64(burst),
		refill:       refill,
		warnInterval: warnInterval,
		now:          now,
		buckets:      make(map[int64]*bucket),
	}
}

// Allow takes one token for key and reports what the caller should do.
func (l *Limiter) Allow(key int64) Decision {
	now := l.now()

	l.mu.Lock()
	defer l.mu.Unlock()

	l.sweep(now)

	b, ok := l.buckets[key]
	if !ok {
		b = &bucket{tokens: l.burst}
		l.buckets[key] = b
	} else if l.refill > 0 {
		b.tokens += float64(now.Sub(b.lastSeen)) / float64(l.refill)
		if b.tokens > l.burst {
			b.tokens = l.burst
		}
	}
	b.lastSeen = now

	if b.tokens >= 1 {
		b.tokens--
		return Decision{Allowed: true}
	}

	decision := Decision{RetryAfter: time.Duration((1 - b.tokens) * float64(l.refill))}
	if b.lastWarn.IsZero() || now.Sub(b.lastWarn) >= l.warnInterval {
		b.lastWarn = now
		decision.Warn = true
	}
	return decision
}

func (l *Limiter) sweep(now time.Time) {
	if l.lastSweep.IsZero() {
		l.lastSweep = now
		return
	}
	if now.Sub(l.lastSweep) < sweepEvery {
		return
	}
	l.lastSweep = now
	for key, b := range l.buckets {
		if now.Sub(b.lastSeen) >= bucketIdleTTL {
			delete(l.buckets, key)
		}
	}
}