// Package ratelimit implements a fixed-window counter per key.
package ratelimit

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

type bucket struct {
	count   int
	resetAt time.Time
}

// Limiter allows at most max events per key within a window.
// The window starts at the first event and is not extended by later ones.
type Limiter struct {
	max    int
	window time.Duration

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

const sweepInterval = time.Minute

func New(max int, window time.Duration) *Limiter {
	return &Limiter{max: max, window: window, buckets: map[string]*bucket{}}
}

// Allow records an event for key and reports whether it is within the limit.
func (l *Limiter) Allow(key string) bool {
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()

	// Prune expired buckets so they do not accumulate. Sweeping on call keeps
	// the package free of goroutines.
	if now.After(l.nextSweep) {
		for k, b := range l.buckets {
			if now.After(b.resetAt) {
				delete(l.buckets, k)
			}
		}
		l.nextSweep = now.Add(sweepInterval)
	}

	b := l.buckets[key]
	if b == nil || now.After(b.resetAt) {
		l.buckets[key] = &bucket{count: 1, resetAt: now.Add(l.window)}
		return true
	}
	if b.count >= l.max {
		return false
	}
	b.count++
	return true
}

// ClientIP returns the address to rate-limit on.
// With a trusted proxy it takes the first X-Forwarded-For entry, because the
// proxy appends the real client there. Otherwise the header is attacker
// controlled and only the socket address can be believed.
func ClientIP(r *http.Request, trustedProxy bool) string {
	if trustedProxy {
		first, _, _ := strings.Cut(r.Header.Get("X-Forwarded-For"), ",")
		return strings.TrimSpace(first)
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}