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

MAX_FAILURES = 10
ACCOUNT_MAX_FAILURES = 5
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)
_account_blocked: dict[str, float] = {}


def _cleanup():
    now = time.time()
    expired_ips = [ip for ip, times in _login_attempts.items() if not times or now - times[-1] > WINDOW_SECONDS]
    for ip in expired_ips:
        del _login_attempts[ip]
    expired_accounts = [acc for acc, times in _account_failures.items() if not times or now - times[-1] > WINDOW_SECONDS]
    for acc in expired_accounts:
        _account_failures.pop(acc, None)
        _account_notified.pop(acc, None)
    expired_blocked = [acc for acc, until in _account_blocked.items() if now > until]
    for acc in expired_blocked:
        del _account_blocked[acc]


def record_login_failure(ip: str, account: str = ""):
    _cleanup()
    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)
        count = len(_account_failures[account])
        if count >= ACCOUNT_MAX_FAILURES:
            _account_blocked[account] = now + WINDOW_SECONDS


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)
        _account_blocked.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_remaining_seconds(account: str) -> int:
    now = time.time()
    if account in _account_blocked:
        remaining = int(_account_blocked[account] - now)
        if remaining > 0:
            return remaining
        del _account_blocked[account]
    return 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


def is_account_blocked(account: str) -> bool:
    return get_account_remaining_seconds(account) > 0


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