package security

import (
	"sync"
	"time"
)

// LoginAttempt 登录尝试记录
type LoginAttempt struct {
	Count       int       // 尝试次数
	LastTry     time.Time // 最后尝试时间
	LockedUntil time.Time // 锁定截止时间
}

// RateLimiter 登录限流器，基于 IP 防止暴力破解
type RateLimiter struct {
	mu              sync.Mutex           // 互斥锁
	attempts        map[string]*LoginAttempt // 尝试记录映射
	maxAttempts     int                  // 最大尝试次数
	lockoutDuration time.Duration        // 锁定时长
}

// NewRateLimiter 创建登录限流器
// @param maxAttempts 最大尝试次数
// @param lockoutDuration 锁定时长（分钟）
// @return *RateLimiter 限流器实例
// @author sunct
func NewRateLimiter(maxAttempts int, lockoutDuration int) *RateLimiter {
	return &RateLimiter{
		attempts:        make(map[string]*LoginAttempt),
		maxAttempts:     maxAttempts,
		lockoutDuration: time.Duration(lockoutDuration) * time.Minute,
	}
}

// Allow 检查 IP 是否允许登录
// @param ip 客户端IP
// @return bool 是否允许
// @return int 剩余尝试次数
// @author sunct
func (rl *RateLimiter) Allow(ip string) (bool, int) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	
	if attempt, exists := rl.attempts[ip]; exists {
		if now.Before(attempt.LockedUntil) {
			return false, rl.maxAttempts - attempt.Count
		}

		if now.Sub(attempt.LastTry) > 1*time.Hour {
			delete(rl.attempts, ip)
			return true, rl.maxAttempts
		}

		if attempt.Count >= rl.maxAttempts {
			attempt.LockedUntil = now.Add(rl.lockoutDuration)
			return false, 0
		}
	}

	return true, rl.maxAttempts
}

// RecordAttempt 记录一次登录尝试
// @param ip 客户端IP
// @author sunct
func (rl *RateLimiter) RecordAttempt(ip string) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	
	if attempt, exists := rl.attempts[ip]; exists {
		attempt.Count++
		attempt.LastTry = now
	} else {
		rl.attempts[ip] = &LoginAttempt{
			Count:   1,
			LastTry: now,
		}
	}
}

// Reset 重置 IP 的登录尝试记录
// @param ip 客户端IP
// @author sunct
func (rl *RateLimiter) Reset(ip string) {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	delete(rl.attempts, ip)
}