package server

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

// RateLimiter is a fixed-window counter keyed by an arbitrary string (client
// address, tunnel ID). A nil limiter allows everything.
type RateLimiter struct {
	mu      sync.Mutex
	limit   int
	window  time.Duration
	entries map[string]*rateEntry
	now     func() time.Time
}

type rateEntry struct {
	count int
	reset time.Time
}

// NewRateLimiter returns a limiter allowing limit events per window, or nil
// when limit is not positive.
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	if limit <= 0 || window <= 0 {
		return nil
	}
	return &RateLimiter{
		limit:   limit,
		window:  window,
		entries: map[string]*rateEntry{},
		now:     time.Now,
	}
}

// Allow records one event for key. When it returns false, retryAfter is the
// time left in the current window.
func (r *RateLimiter) Allow(key string) (ok bool, retryAfter time.Duration) {
	if r == nil {
		return true, 0
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	entry := r.entries[key]
	if entry == nil || !now.Before(entry.reset) {
		if len(r.entries) > 4096 {
			// Keep the map bounded on long-running servers; expired entries
			// are safe to drop.
			for k, e := range r.entries {
				if !now.Before(e.reset) {
					delete(r.entries, k)
				}
			}
		}
		entry = &rateEntry{reset: now.Add(r.window)}
		r.entries[key] = entry
	}
	if entry.count >= r.limit {
		return false, entry.reset.Sub(now)
	}
	entry.count++
	return true, 0
}

// clientIP returns the visitor address, honoring the reverse proxy headers
// the API sits behind. It is only meaningful for rate limiting; direct clients
// can spoof X-Forwarded-For unless the listener is reachable only through the
// proxy.
func clientIP(r *http.Request) string {
	if real := strings.TrimSpace(r.Header.Get("X-Real-IP")); net.ParseIP(real) != nil {
		return real
	}
	if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
		first := strings.TrimSpace(strings.Split(forwarded, ",")[0])
		if net.ParseIP(first) != nil {
			return first
		}
	}
	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return host
	}
	return r.RemoteAddr
}

// retryAfterSeconds formats a Retry-After value (at least one second).
func retryAfterSeconds(d time.Duration) string {
	seconds := int(math.Ceil(d.Seconds()))
	if seconds < 1 {
		seconds = 1
	}
	return strconv.Itoa(seconds)
}