package server

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

type rateLimitState struct {
	count        int
	resetAt      time.Time
	blockedUntil time.Time
}

type rateLimiter struct {
	mu          sync.Mutex
	states      map[string]*rateLimitState
	maxRequests int
	window      time.Duration
	blockFor    time.Duration
}

func newRateLimiter(maxRequests int, window, blockFor time.Duration) *rateLimiter {
	return &rateLimiter{
		states:      make(map[string]*rateLimitState),
		maxRequests: maxRequests,
		window:      window,
		blockFor:    blockFor,
	}
}

func (limiter *rateLimiter) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !limiter.allow(r) {
			w.Header().Set("Retry-After", strconvSeconds(limiter.blockFor))
			http.Error(w, "Too many requests", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (limiter *rateLimiter) allow(r *http.Request) bool {
	key := clientIP(r) + ":" + r.URL.Path
	now := time.Now()

	limiter.mu.Lock()
	defer limiter.mu.Unlock()

	state, ok := limiter.states[key]
	if ok && now.Before(state.blockedUntil) {
		return false
	}

	if !ok || now.After(state.resetAt) {
		limiter.states[key] = &rateLimitState{
			count:   1,
			resetAt: now.Add(limiter.window),
		}
		return true
	}

	state.count++
	if state.count > limiter.maxRequests {
		state.blockedUntil = now.Add(limiter.blockFor)
		state.resetAt = now.Add(limiter.window)
		state.count = 0
		log.Printf("Rate limit triggered for %s on %s", clientIP(r), r.URL.Path)
		return false
	}

	return true
}

func clientIP(r *http.Request) string {
	if forwardedFor := r.Header.Get("X-Forwarded-For"); forwardedFor != "" {
		parts := strings.Split(forwardedFor, ",")
		if ip := strings.TrimSpace(parts[0]); ip != "" {
			return ip
		}
	}
	if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); realIP != "" {
		return realIP
	}
	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return host
	}
	return r.RemoteAddr
}

func strconvSeconds(duration time.Duration) string {
	seconds := int(duration.Seconds())
	if seconds < 1 {
		seconds = 1
	}
	return strconv.Itoa(seconds)
}