import asyncio
import websockets
import json
import logging
import sys
import base64
import os
import argparse
import time

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class WebSocketFileUploader:
    def __init__(self, websocket_url, user_token, space_id):
        self.websocket_url = websocket_url
        self.user_token = user_token
        self.space_id = space_id
        
    async def upload_file(self, websocket, file_path):
        """上传单个文件"""
        if not os.path.exists(file_path):
            logger.error(f"文件不存在: {file_path}")
            return False
        
        filename = os.path.basename(file_path)
        logger.info(f"正在上传文件: {filename}")
        
        try:
            with open(file_path, 'rb') as f:
                file_content = f.read()
                file_b64 = base64.b64encode(file_content).decode('utf-8')
                
                file_msg = {
                    "type": "file_upload",
                    "filename": filename,
                    "content": file_b64,
                    "space_id": self.space_id,
                    "user_token": self.user_token
                }
                
                await websocket.send(json.dumps(file_msg))
                logger.info(f"已发送文件: {filename}")
                
                # 等待服务器响应
                try:
                    response = await asyncio.wait_for(websocket.recv(), timeout=30.0)
                    data = json.loads(response)
                    
                    if data.get("type") == "upload_success":
                        logger.info(f"✅ 文件 '{filename}' 上传成功!")
                        return True
                    elif data.get("type") == "upload_error":
                        logger.error(f"❌ 上传失败: {data.get('message', '未知错误')}")
                        return False
                    else:
                        logger.warning(f"收到未知响应: {data}")
                        return False
                        
                except asyncio.TimeoutError:
                    logger.error(f"上传文件 {filename} 超时")
                    return False
                    
        except Exception as e:
            logger.error(f"上传文件 {file_path} 时出错: {e}")
            return False
    
    async def upload_directory(self, upload_dir):
        """上传目录中的所有文件"""
        if not os.path.exists(upload_dir):
            logger.error(f"目录不存在: {upload_dir}")
            return 0, 0
        
        logger.info(f"🔍 开始扫描目录: {upload_dir}")
        logger.info(f"📡 WebSocket服务器: {self.websocket_url}")
        logger.info(f"🔑 Space ID: {self.space_id}")
        logger.info("-" * 50)
        
        # 获取所有文件
        all_files = []
        for root, dirs, files in os.walk(upload_dir):
            for file in files:
                file_path = os.path.join(root, file)
                if os.path.isfile(file_path):
                    all_files.append(file_path)
        
        if not all_files:
            logger.info("📁 目录中没有找到任何文件")
            return 0, 0
        
        logger.info(f"📁 找到 {len(all_files)} 个文件，开始上传...")
        
        try:
            async with websockets.connect(
                self.websocket_url, 
                ping_interval=20, 
                ping_timeout=60,
                close_timeout=10
            ) as websocket:
                logger.info("WebSocket连接已建立")
                
                # 发送认证信息
                auth_msg = {
                    "type": "auth",
                    "user_token": self.user_token,
                    "space_id": self.space_id
                }
                await websocket.send(json.dumps(auth_msg))
                
                # 等待认证响应
                try:
                    auth_response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    auth_data = json.loads(auth_response)
                    
                    if auth_data.get("type") != "auth_success":
                        logger.error(f"认证失败: {auth_data.get('message', '未知错误')}")
                        return 0, len(all_files)
                        
                    logger.info("认证成功，开始上传文件")
                    
                except asyncio.TimeoutError:
                    logger.error("认证超时")
                    return 0, len(all_files)
                
                success_count = 0
                failed_count = 0
                
                for file_path in all_files:
                    try:
                        if await self.upload_file(websocket, file_path):
                            success_count += 1
                        else:
                            failed_count += 1
                        # 稍微延迟一下，避免服务器压力过大
                        await asyncio.sleep(0.5)
                    except Exception as e:
                        logger.error(f"上传文件 {file_path} 时发生异常: {e}")
                        failed_count += 1
                
                logger.info("-" * 50)
                logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
                
                if success_count > 0:
                    logger.info("🎉 文件已成功上传到您的网盘!")
                
                return success_count, failed_count
                
        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}")
            return 0, len(all_files)

async def upload_directory_websocket(upload_dir, websocket_url, user_token, space_id):
    """使用WebSocket上传目录中的所有文件"""
    uploader = WebSocketFileUploader(websocket_url, user_token, space_id)
    return await uploader.upload_directory(upload_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebSocket文件上传器 - 扫描并上传指定文件夹中的所有文件")
    parser.add_argument("user_token", help="用户令牌 (API Key)")
    parser.add_argument("space_id", help="Space ID")
    parser.add_argument("--websocket-url", default="ws://127.0.0.1:5001/ws", help="WebSocket服务器地址 (默认: ws://127.0.0.1:5001/ws)")
    parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
    
    args = parser.parse_args()
    
    # 开始WebSocket上传
    asyncio.run(upload_directory_websocket(
        args.upload_dir, 
        args.websocket_url, 
        args.user_token, 
        args.space_id
    ))