package builtin

import (
	"sync"
	"time"

	ctxpkg "github.com/noahlann/nnet/pkg/context"
	routerpkg "github.com/noahlann/nnet/pkg/router"
)

// RateLimiter 限流器
type RateLimiter struct {
	tokens    int
	capacity  int
	rate      time.Duration
	lastCheck time.Time
	mu        sync.Mutex
}

// NewRateLimiter 创建限流器
func NewRateLimiter(capacity int, rate time.Duration) *RateLimiter {
	return &RateLimiter{
		tokens:    capacity,
		capacity:  capacity,
		rate:      rate,
		lastCheck: time.Now(),
	}
}

// Allow 检查是否允许
func (rl *RateLimiter) Allow() bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(rl.lastCheck)

	// 根据时间间隔补充令牌
	if elapsed > rl.rate {
		tokensToAdd := int(elapsed / rl.rate)
		rl.tokens += tokensToAdd
		if rl.tokens > rl.capacity {
			rl.tokens = rl.capacity
		}
		rl.lastCheck = now
	}

	if rl.tokens > 0 {
		rl.tokens--
		return true
	}

	return false
}

// RateLimitMiddleware 限流中间件
func RateLimitMiddleware(limiter *RateLimiter) routerpkg.Handler {
	return func(ctx ctxpkg.Context) error {
		if limiter == nil {
			return nil
		}

		if !limiter.Allow() {
			ctx.Response().WriteString("Rate limit exceeded\n")
			return nil
		}

		return nil
	}
}