CheckInApp/apps/backend/services/registration_manager.py

"""
用户名预占和注册限流管理器
"""

import time
import threading
import logging
from typing import Optional, Dict
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)


class RegistrationManager:
    """用户注册管理器 - 处理用户名预占和注册限流"""

    def __init__(self):
        # 用户名预占记录: {alias: {session_id: str, expire_time: float}}
        self._reserved_aliases: Dict[str, Dict] = {}

        # Cookie 注册限流记录: {cookie_value: expire_time}
        self._registration_cookies: Dict[str, float] = {}

        # 线程锁
        self._lock = threading.RLock()

        # 启动清理线程
        self._start_cleanup_thread()

    def reserve_alias(self, alias: str, session_id: str, timeout_seconds: int = 120) -> bool:
        """
        预占用户名

        Args:
            alias: 用户名
            session_id: 会话 ID
            timeout_seconds: 超时时间（秒），默认 120 秒（2 分钟）

        Returns:
            是否预占成功
        """
        with self._lock:
            current_time = time.time()
            expire_time = current_time + timeout_seconds

            # 检查用户名是否已被预占
            if alias in self._reserved_aliases:
                reservation = self._reserved_aliases[alias]

                # 检查是否过期
                if reservation["expire_time"] > current_time:
                    # 未过期，检查是否是同一个 session
                    if reservation["session_id"] == session_id:
                        # 同一个 session，更新过期时间
                        reservation["expire_time"] = expire_time
                        logger.info(f"用户名 {alias} 预占时间已更新（session: {session_id}）")
                        return True
                    else:
                        # 不同 session，预占失败
                        logger.warning(
                            f"用户名 {alias} 已被占用（session: {reservation['session_id']}）"
                        )
                        return False

            # 预占用户名
            self._reserved_aliases[alias] = {"session_id": session_id, "expire_time": expire_time}
            logger.info(f"用户名 {alias} 已预占（session: {session_id}, 超时: {timeout_seconds}s）")
            return True

    def release_alias(self, alias: str, session_id: Optional[str] = None) -> bool:
        """
        释放用户名预占

        Args:
            alias: 用户名
            session_id: 会话 ID（可选，如果提供则只释放匹配的 session）

        Returns:
            是否释放成功
        """
        with self._lock:
            if alias not in self._reserved_aliases:
                return False

            reservation = self._reserved_aliases[alias]

            # 如果指定了 session_id，则只释放匹配的
            if session_id and reservation["session_id"] != session_id:
                logger.warning(f"尝试释放用户名 {alias}，但 session 不匹配")
                return False

            del self._reserved_aliases[alias]
            logger.info(f"用户名 {alias} 预占已释放")
            return True

    def is_alias_reserved(self, alias: str) -> bool:
        """
        检查用户名是否被预占

        Args:
            alias: 用户名

        Returns:
            是否被预占
        """
        with self._lock:
            if alias not in self._reserved_aliases:
                return False

            reservation = self._reserved_aliases[alias]
            current_time = time.time()

            # 检查是否过期
            if reservation["expire_time"] <= current_time:
                # 已过期，自动释放
                del self._reserved_aliases[alias]
                return False

            return True

    def check_registration_cookie(self, cookie_value: str) -> bool:
        """
        检查 Cookie 是否在限流期内

        Args:
            cookie_value: Cookie 值

        Returns:
            True 表示可以注册，False 表示在限流期内
        """
        with self._lock:
            current_time = time.time()

            # 检查 Cookie 是否存在
            if cookie_value in self._registration_cookies:
                expire_time = self._registration_cookies[cookie_value]

                # 检查是否过期
                if expire_time > current_time:
                    remaining = int(expire_time - current_time)
                    logger.warning(
                        f"Cookie {cookie_value[:8]}... 在限流期内（剩余 {remaining} 秒）"
                    )
                    return False
                else:
                    # 已过期，移除记录
                    del self._registration_cookies[cookie_value]

            return True

    def record_registration(self, cookie_value: str, cooldown_seconds: int = 600) -> None:
        """
        记录注册操作（10 分钟冷却）

        Args:
            cookie_value: Cookie 值
            cooldown_seconds: 冷却时间（秒），默认 600 秒（10 分钟）
        """
        with self._lock:
            current_time = time.time()
            expire_time = current_time + cooldown_seconds

            self._registration_cookies[cookie_value] = expire_time
            logger.info(f"Cookie {cookie_value[:8]}... 已记录注册（冷却 {cooldown_seconds} 秒）")

    def _cleanup_expired_records(self) -> None:
        """清理过期的预占记录和限流记录"""
        with self._lock:
            current_time = time.time()

            # 清理过期的用户名预占
            expired_aliases = [
                alias
                for alias, reservation in self._reserved_aliases.items()
                if reservation["expire_time"] <= current_time
            ]

            for alias in expired_aliases:
                del self._reserved_aliases[alias]
                logger.debug(f"用户名 {alias} 预占已过期，自动释放")

            # 清理过期的注册限流记录
            expired_cookies = [
                cookie
                for cookie, expire_time in self._registration_cookies.items()
                if expire_time <= current_time
            ]

            for cookie in expired_cookies:
                del self._registration_cookies[cookie]
                logger.debug(f"Cookie {cookie[:8]}... 限流记录已过期，自动清理")

            if expired_aliases or expired_cookies:
                logger.info(
                    f"清理完成：{len(expired_aliases)} 个用户名，{len(expired_cookies)} 个 Cookie"
                )

    def _start_cleanup_thread(self) -> None:
        """启动定期清理线程"""

        def cleanup_loop():
            while True:
                try:
                    time.sleep(60)  # 每 60 秒清理一次
                    self._cleanup_expired_records()
                except Exception as e:
                    logger.error(f"清理线程异常: {e}")

        thread = threading.Thread(target=cleanup_loop, daemon=True)
        thread.start()
        logger.info("注册管理器清理线程已启动")

    def get_stats(self) -> Dict:
        """获取当前状态统计"""
        with self._lock:
            return {
                "reserved_aliases_count": len(self._reserved_aliases),
                "rate_limited_cookies_count": len(self._registration_cookies),
                "reserved_aliases": list(self._reserved_aliases.keys()),
            }


# 全局单例
registration_manager = RegistrationManager()