package ratelimit

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

	"aijianzhan/platform/internal/authx"
)

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

type visitor struct {
	count int
	reset time.Time
}

func New(ratePerMinute int) *Limiter {
	if ratePerMinute <= 0 {
		ratePerMinute = 120
	}
	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) {
		key := r.RemoteAddr
		if tid := authx.TenantID(r.Context()); tid > 0 {
			key = "t:" + itoa(tid)
		}
		if !l.allow(key) {
			authx.WriteError(w, http.StatusTooManyRequests, "rate limit exceeded")
			return
		}
		next(w, r)
	}
}

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
}

func itoa(n int64) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	if neg {
		i--
		b[i] = '-'
	}
	return string(b[i:])
}