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

class TokenBlacklist:
    """简单的令牌黑名单（内存存储）"""

    def __init__(self):
        self._blacklisted_tokens: Dict[str, datetime] = {}
        self._lock = threading.Lock()

    def add_token(self, token: str, expires_at: Optional[datetime] = None):
        """添加令牌到黑名单"""
        with self._lock:
            # 如果没有提供过期时间，默认24小时后过期
            if expires_at is None:
                expires_at = datetime.utcnow() + timedelta(hours=24)
            self._blacklisted_tokens[token] = expires_at

    def is_blacklisted(self, token: str) -> bool:
        """检查令牌是否在黑名单中"""
        with self._lock:
            if token not in self._blacklisted_tokens:
                return False

            # 检查令牌是否已过期
            if datetime.utcnow() > self._blacklisted_tokens[token]:
                # 清理过期的令牌
                del self._blacklisted_tokens[token]
                return False

            return True

    def cleanup_expired_tokens(self):
        """清理过期的令牌"""
        with self._lock:
            current_time = datetime.utcnow()
            expired_tokens = [
                token for token, expires_at in self._blacklisted_tokens.items()
                if current_time > expires_at
            ]
            for token in expired_tokens:
                del self._blacklisted_tokens[token]

# 全局黑名单实例
token_blacklist = TokenBlacklist()