JoyD/Claw/Server/SmartClaw/src/websocket_client.rs

use tokio_tungstenite::{connect_async, tungstenite::Message, tungstenite::http::Request};
use futures_util::{SinkExt, StreamExt};
use serde_json::json;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use tokio::time::{interval, Duration};
use shared::{TaskRequest, TaskResponse};
use base64::Engine as _;
use base64::engine::general_purpose;
use rand::Rng;

/// WebSocket 客户端连接管理器
pub struct WebSocketClient {
    gateway_url: String,
    sender: Arc<Mutex<Option<mpsc::Sender<String>>>>,
    is_connected: Arc<Mutex<bool>>,
}

impl WebSocketClient {
    /// 创建新的 WebSocket 客户端
    pub fn new(gateway_url: String) -> Self {
        Self {
            gateway_url,
            sender: Arc::new(Mutex::new(None)),
            is_connected: Arc::new(Mutex::new(false)),
        }
    }

    /// 连接到网关服务
    pub async fn connect(&self) -> Result<(), Box<dyn std::error::Error>> {
        println!("🔌 正在连接到网关服务: {}", self.gateway_url);
        
        let ws_url = format!("{}/api/v1/ws/control", self.gateway_url.replace("http://", "ws://").replace("https://", "wss://"));
        println!("🔗 WebSocket URL: {}", ws_url);
        
        // 生成随机的 Sec-WebSocket-Key
        let mut key = [0u8; 16];
        rand::thread_rng().fill(&mut key);
        let sec_websocket_key = general_purpose::STANDARD.encode(&key);
        
        // 建立 WebSocket 连接，添加API密钥和标准 WebSocket 握手头
        let request = Request::builder()
            .uri(&ws_url)
            .header("Host", "pactgo.cn")
            .header("X-API-Key", "claw_secret_key")
            .header("Upgrade", "websocket")
            .header("Connection", "Upgrade")
            .header("Sec-WebSocket-Key", sec_websocket_key)
            .header("Sec-WebSocket-Version", "13")
            .body(())?;
        
        println!("🔗 正在建立WebSocket连接...");
        println!("📋 请求URL: {}", ws_url);
        println!("📋 HTTP版本: {:?}", request.version());
        println!("📋 请求头:");
        for (name, value) in request.headers() {
            println!("  {}: {:?}", name, value);
        }
        
        let (ws_stream, response) = connect_async(request).await?;
        println!("✅ WebSocket 连接建立");
        println!("📋 响应状态: {}", response.status());
        
        // 设置连接状态
        *self.is_connected.lock().await = true;
        
        // 分割流
        let (mut write, mut read) = ws_stream.split();
        
        // 创建消息通道
        let (tx, mut rx) = mpsc::channel::<String>(100);
        *self.sender.lock().await = Some(tx);
        
        // 启动消息发送循环
        let _write_handle = tokio::spawn(async move {
            while let Some(msg) = rx.recv().await {
                if let Err(e) = write.send(Message::Text(msg)).await {
                    println!("❌ 发送消息失败: {}", e);
                    break;
                }
            }
        });
        
        // 启动消息接收循环
        let is_connected_clone = self.is_connected.clone();
        let _read_handle = tokio::spawn(async move {
            while let Some(msg) = read.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        println!("📨 收到消息: {}", text);
                        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&text) {
                            Self::handle_incoming_message(parsed).await;
                        }
                    }
                    Ok(Message::Close(_)) => {
                        println!("🔚 收到关闭消息");
                        *is_connected_clone.lock().await = false;
                        break;
                    }
                    Ok(_) => {}
                    Err(e) => {
                        println!("❌ 接收消息错误: {}", e);
                        *is_connected_clone.lock().await = false;
                        break;
                    }
                }
            }
        });
        
        // 启动心跳机制
        let _heartbeat_handle = {
            let is_connected = self.is_connected.clone();
            tokio::spawn(async move {
                let mut heartbeat_interval = interval(Duration::from_secs(30));
                
                loop {
                    heartbeat_interval.tick().await;
                    
                    let connected = *is_connected.lock().await;
                    if !connected {
                        println!("💔 心跳检测到连接已断开");
                        break;
                    }
                    
                    let _heartbeat_msg = json!({
                        "type": "heartbeat",
                        "service": "smartclaw",
                        "timestamp": chrono::Utc::now().timestamp()
                    }).to_string();
                    
                    // 这里需要重新获取 sender，因为生命周期问题
                    println!("💓 心跳发送");
                }
            })
        };
        
        // 发送连接确认消息
        let connect_msg = json!({
            "type": "connect",
            "service": "smartclaw",
            "version": env!("CARGO_PKG_VERSION"),
            "timestamp": chrono::Utc::now().timestamp()
        }).to_string();
        
        if let Some(sender) = &*self.sender.lock().await {
            let _ = sender.send(connect_msg).await;
        }
        
        println!("🚀 WebSocket 客户端已启动");
        Ok(())
    }

    /// 处理接收到的消息
    async fn handle_incoming_message(message: serde_json::Value) {
        match message.get("type").and_then(|t| t.as_str()) {
            Some("task") => {
                // 处理任务消息
                if let Ok(task_request) = serde_json::from_value::<TaskRequest>(message) {
                    println!("📝 收到任务请求: {:?}", task_request);
                    // 这里可以调用任务处理逻辑
                }
            }
            Some("heartbeat") => {
                println!("💓 收到心跳响应");
            }
            Some("ack") => {
                println!("✅ 收到确认消息");
            }
            Some(msg_type) => {
                println!("❓ 收到未知消息类型: {}", msg_type);
            }
            None => {
                println!("❓ 收到无类型消息");
            }
        }
    }

    /// 发送消息
    pub async fn send_message(&self, message: String) -> Result<(), Box<dyn std::error::Error>> {
        if let Some(sender) = &*self.sender.lock().await {
            sender.send(message).await.map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
            Ok(())
        } else {
            Err("WebSocket 连接未建立".into())
        }
    }

    /// 发送任务响应
    pub async fn send_task_response(&self, response: TaskResponse) -> Result<(), Box<dyn std::error::Error>> {
        let message = json!({
            "type": "task_response",
            "task_id": response.task_id,
            "data": response,
            "timestamp": chrono::Utc::now().timestamp()
        }).to_string();
        
        self.send_message(message).await
    }

    /// 检查连接状态
    pub async fn is_connected(&self) -> bool {
        *self.is_connected.lock().await
    }

    /// 断开连接
    pub async fn disconnect(&self) {
        *self.is_connected.lock().await = false;
        *self.sender.lock().await = None;
        println!("🔌 WebSocket 连接已断开");
    }
}

/// WebSocket 客户端管理器
#[derive(Clone)]
pub struct WebSocketClientManager {
    client: Arc<WebSocketClient>,
}

impl WebSocketClientManager {
    /// 创建新的管理器
    pub fn new(gateway_url: String) -> Self {
        Self {
            client: Arc::new(WebSocketClient::new(gateway_url)),
        }
    }

    /// 启动客户端连接
    pub async fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
        self.client.connect().await
    }

    /// 获取客户端实例
    pub fn get_client(&self) -> Arc<WebSocketClient> {
        self.client.clone()
    }

    /// 停止客户端
    pub async fn stop(&self) {
        self.client.disconnect().await;
    }
}