package shared

import (
	"sync"
	"time"
)

// AttemptLimiter 按 key 统计「失败次数」，达到阈值后在 lock 时长内锁定该 key。
//
// 适用场景：密码、验证码这类可被穷举的入口。成功的尝试应调用 Reset 清零，
// 失败的尝试调用 Fail；处于锁定状态时应直接拒绝，不要再比对密码 —— 既省掉
// 一次密码哈希的计算，也不给攻击者「比对上更快的探测通道」。
//
// 实现约束：
//   - 纯内存、无后台 goroutine：仓库里已经踩过「每个限流器实例一个常驻 goroutine」
//     的泄漏坑（见 internal/route/ratelimit.go 的统一清理循环），这里改用惰性清理；
//   - 键数量上限保护：恶意构造海量 key 时先清过期项，仍超限则丢弃最旧的键，
//     避免 map 无限膨胀把进程内存吃满。
type AttemptLimiter struct {
	mu     sync.Mutex
	fails  map[string][]time.Time
	max    int
	window time.Duration
	lock   time.Duration
	now    func() time.Time // 仅测试替换为可控时钟
}

// NewAttemptLimiter 创建失败计数器：window 内失败达到 max 次 → 锁定 lock 时长。
func NewAttemptLimiter(max int, window, lock time.Duration) *AttemptLimiter {
	if max < 1 {
		max = 1
	}
	return &AttemptLimiter{
		fails:  make(map[string][]time.Time),
		max:    max,
		window: window,
		lock:   lock,
		now:    time.Now,
	}
}

// Blocked 判断 key 是否处于锁定状态，并返回剩余锁定时间。
//
// 锁定期满后自动清零，用户重新获得 max 次机会 —— 否则一次误触发就会永久锁死。
func (l *AttemptLimiter) Blocked(key string) (bool, time.Duration) {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := l.now()
	times := l.pruneLocked(key, now)
	if len(times) < l.max {
		return false, 0
	}

	unlockAt := times[len(times)-1].Add(l.lock)
	if now.Before(unlockAt) {
		return true, unlockAt.Sub(now)
	}

	// 锁定期满：清空记录，重新开始计数。
	delete(l.fails, key)
	return false, 0
}

// Fail 记录一次失败。
func (l *AttemptLimiter) Fail(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := l.now()
	times := l.pruneLocked(key, now)
	times = append(times, now)
	l.fails[key] = times

	// 条数超限时清理全部已过期的键；仍超限则整表重建（宁可丢失计数，也不能 OOM）。
	if len(l.fails) > 5000 {
		for k := range l.fails {
			if t := l.fails[k]; len(t) == 0 || now.Sub(t[len(t)-1]) > l.window+l.lock {
				delete(l.fails, k)
			}
		}
		if len(l.fails) > 5000 {
			l.fails = map[string][]time.Time{key: times}
		}
	}
}

// Reset 清空该 key 的失败记录（校验成功后调用）。
func (l *AttemptLimiter) Reset(key string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	delete(l.fails, key)
}

// pruneLocked 丢弃窗口外的失败记录，返回窗口内的切片（调用方持有 l.mu）。
func (l *AttemptLimiter) pruneLocked(key string, now time.Time) []time.Time {
	times, ok := l.fails[key]
	if !ok {
		return nil
	}
	cutoff := now.Add(-l.window)
	valid := times[:0]
	for _, t := range times {
		if t.After(cutoff) {
			valid = append(valid, t)
		}
	}
	if len(valid) == 0 {
		delete(l.fails, key)
		return nil
	}
	l.fails[key] = valid
	return valid
}