package main

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

// Per-IP token bucket rate limiting (#2). Goroutine-safe via mutex.

type bucket struct {
	tokens float64
	last   time.Time
	rate   float64 // tokens per second
	burst  float64
}

type limiter struct {
	mu      sync.Mutex
	buckets map[string]*bucket
}

func newLimiter() *limiter {
	return &limiter{buckets: make(map[string]*bucket)}
}

func (l *limiter) allow(key string, rate, burst float64) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := time.Now()
	b, ok := l.buckets[key]
	if !ok {
		b = &bucket{tokens: burst, last: now, rate: rate, burst: burst}
		l.buckets[key] = b
	}
	elapsed := now.Sub(b.last).Seconds()
	b.tokens += elapsed * b.rate
	if b.tokens > b.burst {
		b.tokens = b.burst
	}
	b.last = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

// clientIP extracts the request IP (no reverse proxy header by default).
func clientIP(r *http.Request) string {
	host := r.RemoteAddr
	if i := strings.LastIndex(host, ":"); i > 0 {
		host = host[:i]
	}
	return host
}

var globalLimiter = newLimiter()

// rateLimitCreate: 1 req/sec refill, burst 5, per IP.
func rateLimitCreate(r *http.Request) bool {
	return globalLimiter.allow("create:"+clientIP(r), 1, 5)
}

// rateLimitGuess: 1 req/sec refill, burst 5, per IP.
func rateLimitGuess(r *http.Request) bool {
	return globalLimiter.allow("guess:"+clientIP(r), 1, 5)
}

// rateLimitUnlock: 5 per minute per IP+paste.
func rateLimitUnlock(id string, r *http.Request) bool {
	return globalLimiter.allow("unlock:"+id+":"+clientIP(r), 5.0/60.0, 5)
}

// writeRateLimited responds 429 with Retry-After based on refill rate.
func writeRateLimited(w http.ResponseWriter, retryAfterSecs int) {
	w.Header().Set("Retry-After", strconv.Itoa(retryAfterSecs))
	writeErr(w, 429, "rate limit exceeded")
}

// setRateLimitHeaders sets informational X-RateLimit headers for create/guess.
func setRateLimitHeaders(w http.ResponseWriter, limit, burst int) {
	w.Header().Set("X-RateLimit-Limit", strconv.Itoa(limit))
	w.Header().Set("X-RateLimit-Burst", strconv.Itoa(burst))
}