import time
from collections import defaultdict
from fastapi import Request, HTTPException

MAX_FAILURES = 10
WINDOW_SECONDS = 300
ALERT_THRESHOLDS = [3, 5, 10]

_login_attempts: dict[str, list[float]] = defaultdict(list)
_account_failures: dict[str, list[float]] = defaultdict(list)
_account_notified: dict[str, set[int]] = defaultdict(set)


def record_login_failure(ip: str, account: str = ""):
    now = time.time()
    _login_attempts[ip] = [t for t in _login_attempts[ip] if now - t < WINDOW_SECONDS]
    _login_attempts[ip].append(now)
    if account:
        _account_failures[account] = [t for t in _account_failures[account] if now - t < WINDOW_SECONDS]
        _account_failures[account].append(now)


def clear_login_failures(ip: str, account: str = ""):
    _login_attempts.pop(ip, None)
    if account:
        _account_failures.pop(account, None)
        _account_notified.pop(account, None)


def get_remaining_seconds(ip: str) -> int:
    now = time.time()
    _login_attempts[ip] = [t for t in _login_attempts[ip] if now - t < WINDOW_SECONDS]
    if not _login_attempts[ip]:
        return 0
    oldest = _login_attempts[ip][0]
    remaining = int(WINDOW_SECONDS - (now - oldest))
    return max(remaining, 0)


def get_account_failure_count(account: str) -> int:
    now = time.time()
    _account_failures[account] = [t for t in _account_failures[account] if now - t < WINDOW_SECONDS]
    return len(_account_failures[account])


def get_unnotified_threshold(account: str) -> int | None:
    now = time.time()
    _account_failures[account] = [t for t in _account_failures[account] if now - t < WINDOW_SECONDS]
    count = len(_account_failures[account])
    notified = _account_notified[account]
    for threshold in ALERT_THRESHOLDS:
        if count >= threshold and threshold not in notified:
            notified.add(threshold)
            return threshold
    return None


def is_ip_blocked(ip: str) -> bool:
    return get_remaining_seconds(ip) > 0 and len(_login_attempts[ip]) >= MAX_FAILURES


async def ip_rate_limit_middleware(request: Request, call_next):
    ip = request.client.host if request.client else "unknown"
    if is_ip_blocked(ip):
        remaining = get_remaining_seconds(ip)
        raise HTTPException(status_code=429, detail=f"登录失败次数过多，请 {remaining} 秒后再试")
    response = await call_next(request)
    return response