cvkg-cli 0.1.6

//! WebSocket Server
//! Multiplexed WebSocket server for runtime communication, DevTools, hot reload, and agent streams

use axum::{
    extract::ws::{WebSocketUpgrade, WebSocket, Message},
    response::IntoResponse,   routing::get,
    Router,
};
use futures_util::StreamExt;
use std::net::SocketAddr;
use tracing::{info, error};

use serde::{Deserialize, Serialize};

/// WebSocket message types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum WsMessage {
    Patch(super::patch_engine::RuntimePatch),
    State(super::dev_runtime::RuntimeStateSnapshot),
    Event(super::dev_runtime::RuntimeEvent),
    Devtools(serde_json::Value),
}

/// WebSocket handler for runtime communication
async fn runtime_ws(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_runtime_socket)
}

/// WebSocket handler for DevTools
async fn devtools_ws(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_devtools_socket)
}

/// WebSocket handler for hot reload
async fn hotreload_ws(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_hotreload_socket)
}

/// WebSocket handler for agent streams
async fn agent_ws(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_agent_socket)
}

/// Handle runtime WebSocket connection
async fn handle_runtime_socket(mut ws: WebSocket) {
    info!("Runtime WebSocket client connected");
    
    // Send initial handshake
    let _ = ws
        .send(Message::Text(
            serde_json::to_string(&serde_json::json!({
                "type": "handshake",
                "payload": {
                    "client": "runtime",
                    "capabilities": ["patch", "state", "event"]
                }
            })).unwrap().into(),
        ))
        .await;
    
    while let Some(result) = ws.next().await {
        match result {
            Ok(Message::Text(text)) => {
                // Handle incoming runtime messages
                if let Ok(message) = serde_json::from_str::<serde_json::Value>(&text) {
                    // Process runtime messages
                    info!("Received runtime message: {}", message);
                }
            }
            Ok(Message::Binary(bin)) => {
                // Handle binary messages if needed
                info!("Received binary message of {} bytes", bin.len());
            }
            Ok(Message::Close(_)) => {
                info!("Runtime WebSocket client disconnected");
                break;
            }
            Err(e) => {
                error!("WebSocket error: {}", e);
                break;
            }
            _ => {}
        }
    }
}

/// Handle DevTools WebSocket connection
async fn handle_devtools_socket(mut ws: WebSocket) {
    info!("DevTools WebSocket client connected");
    
    while let Some(result) = ws.next().await {
        match result {
            Ok(Message::Text(text)) => {
                // Handle DevTools messages
                if let Ok(message) = serde_json::from_str::<serde_json::Value>(&text) {
                    info!("Received DevTools message: {}", message);
                }
            }
            Ok(Message::Close(_)) => {
                info!("DevTools WebSocket client disconnected");
                break;
            }
            Err(e) => {
                error!("DevTools WebSocket error: {}", e);
                break;
            }
            _ => {}
        }
    }
}

/// Handle hot reload WebSocket connection
async fn handle_hotreload_socket(mut ws: WebSocket) {
    info!("Hot reload WebSocket client connected");
    
    while let Some(result) = ws.next().await {
        match result {
            Ok(Message::Text(text)) => {
                // Handle hot reload messages (patch delivery)
                if let Ok(message) = serde_json::from_str::<WsMessage>(&text) {
                    if let WsMessage::Patch(patch) = message {
                        info!("Received patch for hot reload: {:?}", patch);
                        // TODO: Apply patch to runtime
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Hot reload WebSocket client disconnected");
                break;
            }
            Err(e) => {
                error!("Hot reload WebSocket error: {}", e);
                break;
            }
            _ => {}
        }
    }
}

/// Handle agent stream WebSocket connection
async fn handle_agent_socket(mut ws: WebSocket) {
    info!("Agent stream WebSocket client connected");
    
    while let Some(result) = ws.next().await {
        match result {
            Ok(Message::Text(text)) => {
                // Handle agent stream messages
                if let Ok(message) = serde_json::from_str::<serde_json::Value>(&text) {
                    info!("Received agent message: {}", message);
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent stream WebSocket client disconnected");
                break;
            }
            Err(e) => {
                error!("Agent stream WebSocket error: {}", e);
                break;
            }
            _ => {}
        }
    }
}

/// Create the WebSocket router with all endpoints
pub fn create_router() -> Router {
    Router::new()
        .route("/ws/runtime", get(runtime_ws))
        .route("/ws/devtools", get(devtools_ws))
        .route("/ws/hotreload", get(hotreload_ws))
        .route("/ws/agent", get(agent_ws))
}

/// Start the WebSocket server
pub async fn start_server(addr: SocketAddr) -> Result<(), Box<dyn std::error::Error>> {
    let app = create_router();
    info!("Starting WebSocket server on {}", addr);
    
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;
    
    Ok(())
}