wavecraft_dev_server/ws/

mod.rs

//! WebSocket server for browser-based UI development
//!
//! This module provides a WebSocket server that exposes the same IPC protocol
//! used by the native WKWebView transport, enabling real-time communication
//! between a browser-based UI and the Rust engine during development.

use futures_util::{SinkExt, StreamExt};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{RwLock, broadcast};
use tokio_tungstenite::{accept_async, tungstenite::protocol::Message};
use tracing::{debug, error, info, warn};
use wavecraft_bridge::{IpcHandler, ParameterHost};

/// Shared state for tracking connected clients
struct ServerState {
    /// Connected browser clients (for broadcasting meter updates)
    browser_clients: Arc<RwLock<Vec<tokio::sync::mpsc::Sender<String>>>>,
    /// Audio client ID (if connected)
    audio_client: Arc<RwLock<Option<String>>>,
}

impl ServerState {
    fn new() -> Self {
        Self {
            browser_clients: Arc::new(RwLock::new(Vec::new())),
            audio_client: Arc::new(RwLock::new(None)),
        }
    }
}

/// A lightweight, cloneable handle to the WebSocket server's broadcast
/// capability. Non-generic — can be passed across async task boundaries.
///
/// Constructed via [`WsServer::handle()`]. Used by the CLI to forward
/// meter updates from the in-process audio callback to browser clients.
#[derive(Clone)]
pub struct WsHandle {
    state: Arc<ServerState>,
}

impl WsHandle {
    /// Broadcast a JSON string to all connected browser clients.
    pub async fn broadcast(&self, json: &str) {
        let clients = self.state.browser_clients.read().await;
        for client in clients.iter() {
            if let Err(e) = client.try_send(json.to_owned()) {
                warn!("Failed to broadcast message to client: {}", e);
            }
        }
    }
}

/// WebSocket server for browser-based UI development
pub struct WsServer<H: ParameterHost + 'static> {
    /// Port the server listens on
    port: u16,
    /// Shared IPC handler
    handler: Arc<IpcHandler<H>>,
    /// Shutdown signal
    shutdown_tx: broadcast::Sender<()>,
    /// Enable verbose logging (all JSON-RPC messages)
    verbose: bool,
    /// Shared server state
    state: Arc<ServerState>,
}

impl<H: ParameterHost + 'static> WsServer<H> {
    /// Create a new WebSocket server
    pub fn new(port: u16, handler: Arc<IpcHandler<H>>, verbose: bool) -> Self {
        let (shutdown_tx, _) = broadcast::channel(1);
        Self {
            port,
            handler,
            shutdown_tx,
            verbose,
            state: Arc::new(ServerState::new()),
        }
    }

    /// Get a lightweight handle for broadcasting to connected clients.
    ///
    /// The returned `WsHandle` is non-generic, `Clone`, and can be moved
    /// into async tasks (e.g., for forwarding meter updates from audio).
    pub fn handle(&self) -> WsHandle {
        WsHandle {
            state: Arc::clone(&self.state),
        }
    }

    /// Broadcast a parametersChanged notification to all connected clients.
    ///
    /// This is used by the hot-reload pipeline to notify the UI that
    /// parameters have been updated and should be re-fetched.
    pub async fn broadcast_parameters_changed(&self) -> Result<(), serde_json::Error> {
        use wavecraft_protocol::IpcNotification;

        let notification = IpcNotification::new("parametersChanged", serde_json::json!({}));
        let json = serde_json::to_string(&notification)?;

        let clients = self.state.browser_clients.read().await;
        for client in clients.iter() {
            if let Err(e) = client.try_send(json.clone()) {
                warn!("Failed to send parametersChanged notification to client: {}", e);
            }
        }

        Ok(())
    }

    /// Start the server (spawns async tasks, returns immediately)
    pub async fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
        let addr: SocketAddr = format!("127.0.0.1:{}", self.port).parse()?;
        let listener = TcpListener::bind(&addr).await?;

        info!("Server listening on ws://{}", addr);

        let handler = Arc::clone(&self.handler);
        let mut shutdown_rx = self.shutdown_tx.subscribe();
        let verbose = self.verbose;
        let state = Arc::clone(&self.state);

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    result = listener.accept() => {
                        match result {
                            Ok((stream, addr)) => {
                                info!("Client connected: {}", addr);
                                let handler = Arc::clone(&handler);
                                let state = Arc::clone(&state);
                                tokio::spawn(handle_connection(handler, stream, addr, verbose, state));
                            }
                            Err(e) => {
                                error!("Accept error: {}", e);
                            }
                        }
                    }
                    _ = shutdown_rx.recv() => {
                        info!("Server shutting down");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// Shutdown the server gracefully.
    ///
    /// Note: Not currently called but kept for future graceful shutdown support.
    #[allow(dead_code)]
    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(());
    }
}

/// Handle a single WebSocket connection
async fn handle_connection<H: ParameterHost>(
    handler: Arc<IpcHandler<H>>,
    stream: TcpStream,
    addr: SocketAddr,
    verbose: bool,
    state: Arc<ServerState>,
) {
    let ws_stream = match accept_async(stream).await {
        Ok(ws) => ws,
        Err(e) => {
            error!("Error during handshake with {}: {}", addr, e);
            return;
        }
    };

    info!("WebSocket connection established: {}", addr);

    let (mut write, mut read) = ws_stream.split();
    let (tx, mut rx) = tokio::sync::mpsc::channel::<String>(128);

    // Track this client for broadcasting
    let mut is_audio_client = false;
    let client_index = {
        let mut clients = state.browser_clients.write().await;
        clients.push(tx.clone());
        clients.len() - 1
    };

    // Spawn task to send messages from channel to WebSocket
    let write_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if let Err(e) = write.send(Message::Text(msg)).await {
                error!("Error sending to {}: {}", addr, e);
                break;
            }
        }
    });

    while let Some(msg) = read.next().await {
        match msg {
            Ok(Message::Text(json)) => {
                // Log incoming message (verbose only)
                if verbose {
                    debug!("Received from {}: {}", addr, json);
                }

                // Try to parse as IPC request for structured routing
                let parsed_req = serde_json::from_str::<wavecraft_protocol::IpcRequest>(&json);
                
                if let Ok(ref req) = parsed_req {
                    // Handle registerAudio
                    if req.method == "registerAudio" {
                        is_audio_client = true;
                        info!("Audio client registered: {}", addr);

                        // Parse to extract client_id
                        if let Some(params) = req.params.clone()
                            && let Ok(audio_params) = serde_json::from_value::<
                                wavecraft_protocol::RegisterAudioParams,
                            >(params)
                        {
                            *state.audio_client.write().await = Some(audio_params.client_id.clone());
                        }

                        // Send success response using the request's id
                        let response = wavecraft_protocol::IpcResponse::success(
                            req.id.clone(),
                            wavecraft_protocol::RegisterAudioResult {
                                status: "registered".to_string(),
                            },
                        );
                        let response_json = match serde_json::to_string(&response) {
                            Ok(json) => json,
                            Err(e) => {
                                error!("Failed to serialize registerAudio response: {}", e);
                                break;
                            }
                        };
                        if let Err(e) = tx.try_send(response_json) {
                            error!("Error sending response: {}", e);
                            break;
                        }
                        continue;
                    }
                    
                    // Handle meterUpdate from audio client
                    if is_audio_client && req.method == "meterUpdate" {
                        // Broadcast to all browser clients
                        let clients = state.browser_clients.read().await;
                        for (idx, client) in clients.iter().enumerate() {
                            if idx != client_index {
                                // Don't send back to audio client
                                if let Err(e) = client.try_send(json.clone()) {
                                    warn!("Failed to broadcast meter update to client {}: {}", idx, e);
                                }
                            }
                        }
                        continue;
                    }
                }

                // Route through existing IpcHandler
                let response = handler.handle_json(&json);

                // Log outgoing response (verbose only)
                if verbose {
                    debug!("Sending to {}: {}", addr, response);
                }

                // Send response
                if let Err(e) = tx.try_send(response) {
                    error!("Error queueing response: {}", e);
                    break;
                }
            }
            Ok(Message::Close(_)) => {
                info!("Client closed connection: {}", addr);
                break;
            }
            Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {
                // Ignore ping/pong frames (automatically handled)
            }
            Ok(Message::Binary(_)) => {
                warn!("Unexpected binary message from {}", addr);
            }
            Ok(Message::Frame(_)) => {
                // Raw frames shouldn't appear at this level
            }
            Err(e) => {
                error!("Error receiving from {}: {}", addr, e);
                break;
            }
        }
    }

    // Cleanup: remove client from broadcast list
    state
        .browser_clients
        .write()
        .await
        .retain(|c| !c.is_closed());
    if is_audio_client {
        *state.audio_client.write().await = None;
        info!("Audio client disconnected: {}", addr);
    }

    write_task.abort();
    info!("Connection closed: {}", addr);
}

#[cfg(test)]
mod tests {
    use super::*;
    use wavecraft_bridge::InMemoryParameterHost;
    use wavecraft_protocol::{ParameterInfo, ParameterType};

    /// Simple test host for unit tests
    fn test_host() -> InMemoryParameterHost {
        InMemoryParameterHost::new(vec![ParameterInfo {
            id: "gain".to_string(),
            name: "Gain".to_string(),
            param_type: ParameterType::Float,
            value: 0.5,
            default: 0.5,
            unit: Some("dB".to_string()),
            group: Some("Input".to_string()),
        }])
    }

    #[tokio::test]
    async fn test_server_creation() {
        let host = test_host();
        let handler = Arc::new(IpcHandler::new(host));
        let server = WsServer::new(9001, handler, false);

        // Just verify we can create a server without panicking
        assert_eq!(server.port, 9001);
    }
}