"""
Token认证系统
基于用户信息生成访问令牌，类似HubApi的使用方式
"""
import secrets
import hashlib
import json
import os
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any
from pydantic import BaseModel
from enum import Enum


class UserRole(str, Enum):
    """用户角色"""
    ADMIN = "admin"
    USER = "user"
    GUEST = "guest"


class TokenRequest(BaseModel):
    """获取Token的请求"""
    user_id: str
    user_name: str
    role: UserRole = UserRole.USER
    
    # 权限控制
    max_file_size: int = 1024 * 1024 * 1024 * 1024  # 1TB默认
    daily_upload_limit: int = 300  # 每日上传次数限制


class TokenResponse(BaseModel):
    """Token响应"""
    access_token: str
    token_type: str = "Bearer"
    user_info: Dict[str, Any]


class TokenInfo(BaseModel):
    """Token信息（内部存储）"""
    token_hash: str
    user_id: str
    user_name: str
    role: UserRole
    
    # 权限和限制
    max_file_size: int = 1024 * 1024 * 1024 * 1024  # 1TB
    daily_upload_limit: int = 300
    
    # 时间信息
    created_at: datetime
    last_used_at: Optional[datetime] = None
    
    # 使用统计
    usage_count: int = 0
    daily_usage: int = 0


class AuthenticatedUser(BaseModel):
    """认证后的用户信息"""
    user_id: str
    user_name: str
    role: UserRole
    max_file_size: int
    daily_remaining: int
    token_hash: str


class TokenManager:
    """Token管理器"""
    
    def __init__(self, data_file: str = "tokens.json"):
        self.data_file = data_file
        self.tokens = self._load_tokens()
    
    def _load_tokens(self) -> Dict[str, TokenInfo]:
        """加载Token数据"""
        if os.path.exists(self.data_file):
            try:
                with open(self.data_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    tokens = {}
                    for token_hash, token_data in data.items():
                        # 处理datetime字段反序列化
                        if 'created_at' in token_data and isinstance(token_data['created_at'], str):
                            token_data['created_at'] = datetime.fromisoformat(token_data['created_at'])
                        if 'expires_at' in token_data and isinstance(token_data['expires_at'], str):
                            token_data['expires_at'] = datetime.fromisoformat(token_data['expires_at'])
                        if 'last_used_at' in token_data and token_data['last_used_at'] and isinstance(token_data['last_used_at'], str):
                            token_data['last_used_at'] = datetime.fromisoformat(token_data['last_used_at'])
                        
                        tokens[token_hash] = TokenInfo(**token_data)
                    return tokens
            except Exception as e:
                print(f"加载Token数据失败: {e}")
        return {}
    
    def _save_tokens(self):
        """保存Token数据"""
        data = {}
        for token_hash, token_info in self.tokens.items():
            token_dict = token_info.dict()
            # 处理datetime序列化
            if token_dict['created_at']:
                token_dict['created_at'] = token_dict['created_at'].isoformat()
            if token_dict['expires_at']:
                token_dict['expires_at'] = token_dict['expires_at'].isoformat()
            if token_dict['last_used_at']:
                token_dict['last_used_at'] = token_dict['last_used_at'].isoformat()
            
            data[token_hash] = token_dict
        
        with open(self.data_file, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    def generate_token(self, request: TokenRequest) -> TokenResponse:
        """生成访问令牌"""
        # 生成随机令牌
        token = f"hf_{secrets.token_urlsafe(32)}"  # 类似HuggingFace的格式
        token_hash = hashlib.sha256(token.encode()).hexdigest()
        
        # 计算过期时间
        expires_at = datetime.now() + timedelta(hours=request.expires_hours)
        
        # 创建Token信息
        token_info = TokenInfo(
            token_hash=token_hash,
            user_id=request.user_id,
            user_name=request.user_name,
            role=request.role,
            email=request.email,
            allowed_datasets=request.allowed_datasets,
            max_file_size=request.max_file_size,
            daily_upload_limit=request.daily_upload_limit,
            created_at=datetime.now(),
            expires_at=expires_at,
            last_reset_date=datetime.now().strftime("%Y-%m-%d")
        )
        
        # 保存Token
        self.tokens[token_hash] = token_info
        self._save_tokens()
        
        # 返回响应
        return TokenResponse(
            access_token=token,
            expires_in=int(request.expires_hours * 3600),
            expires_at=expires_at,
            user_info={
                "user_id": request.user_id,
                "user_name": request.user_name,
                "role": request.role.value,
                "email": request.email,
                "allowed_datasets": request.allowed_datasets,
                "max_file_size": request.max_file_size,
                "daily_upload_limit": request.daily_upload_limit
            }
        )
    
    def verify_token(self, token: str) -> Optional[AuthenticatedUser]:
        """验证Token并返回用户信息"""
        if not token or not token.startswith("hf_"):
            return None
        
        token_hash = hashlib.sha256(token.encode()).hexdigest()
        token_info = self.tokens.get(token_hash)
        
        if not token_info:
            return None
        
        # 检查是否过期
        if datetime.now() > token_info.expires_at:
            return None
        
        # 更新使用统计
        self._update_token_usage(token_info)
        
        # 计算今日剩余次数
        daily_remaining = max(0, token_info.daily_upload_limit - token_info.daily_usage)
        
        return AuthenticatedUser(
            user_id=token_info.user_id,
            user_name=token_info.user_name,
            role=token_info.role,
            email=token_info.email,
            allowed_datasets=token_info.allowed_datasets,
            max_file_size=token_info.max_file_size,
            daily_remaining=daily_remaining
        )
    
    def _update_token_usage(self, token_info: TokenInfo):
        """更新Token使用统计"""
        now = datetime.now()
        today = now.strftime("%Y-%m-%d")
        
        # 如果是新的一天，重置每日使用计数
        if token_info.last_reset_date != today:
            token_info.daily_usage = 0
            token_info.last_reset_date = today
        
        # 更新使用统计
        token_info.usage_count += 1
        token_info.daily_usage += 1
        token_info.last_used_at = now
        
        self._save_tokens()
    
    def revoke_token(self, token: str) -> bool:
        """撤销Token"""
        if not token or not token.startswith("hf_"):
            return False
        
        token_hash = hashlib.sha256(token.encode()).hexdigest()
        if token_hash in self.tokens:
            del self.tokens[token_hash]
            self._save_tokens()
            return True
        return False
    
    def get_user_tokens(self, user_id: str) -> List[Dict[str, Any]]:
        """获取用户的所有Token（不包含实际token值）"""
        user_tokens = []
        for token_hash, token_info in self.tokens.items():
            if token_info.user_id == user_id:
                user_tokens.append({
                    "token_hash": token_hash[:16] + "...",  # 只显示部分哈希
                    "user_name": token_info.user_name,
                    "role": token_info.role.value,
                    "created_at": token_info.created_at,
                    "expires_at": token_info.expires_at,
                    "last_used_at": token_info.last_used_at,
                    "usage_count": token_info.usage_count,
                    "daily_usage": token_info.daily_usage,
                    "daily_limit": token_info.daily_upload_limit
                })
        return user_tokens


# 全局Token管理器实例
token_manager = TokenManager()


def get_token_manager() -> TokenManager:
    """获取Token管理器实例"""
    return token_manager
