package main

import (
	"sync"
	"time"
)

// bucket is a token bucket: refill tokens/sec up to burst.
type bucket struct {
	mu     sync.Mutex
	tokens float64
	burst  float64
	rate   float64
	last   time.Time
}

func newBucket(rate, burst float64) *bucket {
	return &bucket{tokens: burst, burst: burst, rate: rate, last: time.Now()}
}

func (b *bucket) take() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	now := time.Now()
	b.tokens += now.Sub(b.last).Seconds() * b.rate
	if b.tokens > b.burst {
		b.tokens = b.burst
	}
	b.last = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

// ipRateLimiter allows n requests per window per IP, with periodic cleanup.
type ipRateLimiter struct {
	mu     sync.Mutex
	n      int
	window time.Duration
	hits   map[string][]time.Time
}

func newIPRateLimiter(n int, window time.Duration) *ipRateLimiter {
	l := &ipRateLimiter{n: n, window: window, hits: map[string][]time.Time{}}
	go func() {
		for range time.Tick(5 * time.Minute) {
			l.mu.Lock()
			cutoff := time.Now().Add(-l.window)
			for ip, ts := range l.hits {
				if len(ts) == 0 || ts[len(ts)-1].Before(cutoff) {
					delete(l.hits, ip)
				}
			}
			l.mu.Unlock()
		}
	}()
	return l
}

func (l *ipRateLimiter) allow(ip string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := time.Now()
	cutoff := now.Add(-l.window)
	ts := l.hits[ip]
	kept := ts[:0]
	for _, t := range ts {
		if t.After(cutoff) {
			kept = append(kept, t)
		}
	}
	if len(kept) >= l.n {
		l.hits[ip] = kept
		return false
	}
	l.hits[ip] = append(kept, now)
	return true
}