import asyncio
import websockets
import subprocess
import json
import logging
import sys
import urllib.parse
import base64
import os
import threading
import requests
from urllib.parse import urlparse
import argparse

# Configure logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 全局变量，用于在用户输入时访问 WebSocket
global_websocket = None

async def connect(card_id, machine_secret, upload_files=None, server="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/"):
    global global_websocket

    # 使用查询参数传递 machine_secret
    encoded_secret = urllib.parse.quote(machine_secret)
    uri = f"{server}{card_id}?secret={encoded_secret}"
    logger.info(f"Attempting to connect to {uri}")

    try:
        async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
            global_websocket = websocket
            logger.info("Connected to WebSocket")

            machine_info = {"type": "machine", "card_id": card_id, "token": machine_secret}  # 添加用户token
            await websocket.send(json.dumps(machine_info))
            logger.debug(f"Sent machine registration: {machine_info}")

            # 如果指定了上传文件列表，发送文件内容
            if upload_files:
                for upload_file in upload_files:
                    if os.path.exists(upload_file):
                        try:
                            with open(upload_file, 'rb') as f:
                                file_content = f.read()
                            file_b64 = base64.b64encode(file_content).decode('utf-8')
                            file_msg = {
                                "type": "file_upload",
                                "filename": os.path.basename(upload_file),
                                "content": file_b64
                                # 去除 machine_id
                            }
                            await websocket.send(json.dumps(file_msg))
                            logger.info(f"Uploaded file: {upload_file}")
                        except Exception as e:
                            logger.error(f"Error uploading file {upload_file}: {e}")
                    else:
                        logger.error(f"File not found: {upload_file}")

            # 等待一段时间以接收任何响应，然后关闭
            try:
                await asyncio.sleep(5)  # 等待服务器响应
            except:
                pass

    except Exception as e:
        logger.error(f"Failed to connect or maintain connection: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="WebSocket-based File Uploader for remote control file transfer")
    parser.add_argument("api_key", help="API Key (used as machine_secret and token)")
    parser.add_argument("space_id", help="Space ID (used as card_id)")
    parser.add_argument("--server", default="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/", help="WebSocket server base URL (default: wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/)")
    parser.add_argument("--upload-dir", default="output", help="Directory to upload files from (default: output)")

    args = parser.parse_args()

    if not os.path.exists(args.upload_dir):
        print(f"目录不存在: {args.upload_dir}")
        sys.exit(1)

    all_files = []
    for root, dirs, files in os.walk(args.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:
        print("目录中没有找到任何文件")
        sys.exit(0)

    print(f"找到 {len(all_files)} 个文件，开始上传...")
    asyncio.run(connect(args.space_id, args.api_key, upload_files=all_files, server=args.server))
    print("上传完成!")