package ratelimit

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

type Limiter struct {
	mu       sync.Mutex
	visitors map[string]*visitor
	rate     int // 0 = disabled
	window   time.Duration
}

type visitor struct {
	count int
	reset time.Time
}

// New creates a per-IP limiter. ratePerMinute <= 0 disables limiting.
func New(ratePerMinute int) *Limiter {
	return &Limiter{visitors: map[string]*visitor{}, rate: ratePerMinute, window: time.Minute}
}

func (l *Limiter) Middleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if l.rate <= 0 {
			next(w, r)
			return
		}
		ip, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil || ip == "" {
			ip = r.RemoteAddr
		}
		// 本机开发：Vite 代理与所有本机请求共用 127.0.0.1，不做限流
		if isLoopback(ip) {
			next(w, r)
			return
		}
		if !l.allow(ip) {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusTooManyRequests)
			_, _ = w.Write([]byte(`{"code":429,"message":"gateway rate limit exceeded"}`))
			return
		}
		next(w, r)
	}
}

func isLoopback(ip string) bool {
	ip = strings.Trim(ip, "[]")
	if ip == "127.0.0.1" || ip == "::1" || ip == "localhost" {
		return true
	}
	parsed := net.ParseIP(ip)
	return parsed != nil && parsed.IsLoopback()
}

func (l *Limiter) allow(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := time.Now()
	v, ok := l.visitors[key]
	if !ok || now.After(v.reset) {
		l.visitors[key] = &visitor{count: 1, reset: now.Add(l.window)}
		return true
	}
	if v.count >= l.rate {
		return false
	}
	v.count++
	return true
}