package ratelimit

import (
	"sync"
	"time"
)

// Limiter 内存限流器
type Limiter struct {
	mu sync.Mutex

	// 每用户每秒请求数
	userRPS map[uint64]*tokenBucket

	// 密钥每日请求计数
	keyDailyReq map[uint64]*dailyCounter

	// 密钥每日 token 计数
	keyDailyTokens map[uint64]*dailyCounter
}

type tokenBucket struct {
	tokens   float64
	maxTokens float64
	refillRate float64
	lastRefill time.Time
}

type dailyCounter struct {
	date  string
	count int64
}

func New() *Limiter {
	return &Limiter{
		userRPS:       make(map[uint64]*tokenBucket),
		keyDailyReq:   make(map[uint64]*dailyCounter),
		keyDailyTokens: make(map[uint64]*dailyCounter),
	}
}

// AllowRequest 检查用户级每秒请求限制
func (l *Limiter) AllowRequest(userID uint64, rps int) bool {
	if rps <= 0 {
		return true
	}
	l.mu.Lock()
	defer l.mu.Unlock()

	bucket, ok := l.userRPS[userID]
	if !ok {
		bucket = &tokenBucket{
			tokens:     float64(rps),
			maxTokens:  float64(rps),
			refillRate: float64(rps),
			lastRefill: time.Now(),
		}
		l.userRPS[userID] = bucket
	}

	now := time.Now()
	elapsed := now.Sub(bucket.lastRefill).Seconds()
	bucket.tokens += elapsed * bucket.refillRate
	if bucket.tokens > bucket.maxTokens {
		bucket.tokens = bucket.maxTokens
	}
	bucket.lastRefill = now

	if bucket.tokens < 1 {
		return false
	}
	bucket.tokens--
	return true
}

// AllowRequestDaily 检查密钥每日请求配额
func (l *Limiter) AllowRequestDaily(keyID uint64, quota int) bool {
	if quota <= 0 {
		return true
	}
	l.mu.Lock()
	defer l.mu.Unlock()

	today := time.Now().UTC().Format("2006-01-02")
	counter, ok := l.keyDailyReq[keyID]
	if !ok || counter.date != today {
		l.keyDailyReq[keyID] = &dailyCounter{date: today, count: 1}
		return true
	}
	if counter.count >= int64(quota) {
		return false
	}
	counter.count++
	return true
}

// TokensUsed 返回密钥今日 token 用量
func (l *Limiter) TokensUsed(keyID uint64) int64 {
	l.mu.Lock()
	defer l.mu.Unlock()

	today := time.Now().UTC().Format("2006-01-02")
	counter, ok := l.keyDailyTokens[keyID]
	if !ok || counter.date != today {
		return 0
	}
	return counter.count
}

// AddTokens 累加密钥今日 token 用量
func (l *Limiter) AddTokens(keyID uint64, tokens int64) {
	l.mu.Lock()
	defer l.mu.Unlock()

	today := time.Now().UTC().Format("2006-01-02")
	counter, ok := l.keyDailyTokens[keyID]
	if !ok || counter.date != today {
		l.keyDailyTokens[keyID] = &dailyCounter{date: today, count: tokens}
		return
	}
	counter.count += tokens
}