graphwalker_restful/

websocket.rs

use std::collections::HashSet;
use std::time::Duration;

use axum::extract::ws::{Message, WebSocket};
use futures::{SinkExt, StreamExt};
use serde_json::{json, Value};
use tokio::sync::{broadcast, oneshot};
use tracing::{debug, warn};

use crate::actor::{self, spawn_machine_thread, Command};
use crate::session::{SessionHandle, SessionManager};

pub async fn handle_socket(socket: WebSocket, session_mgr: SessionManager) {
    let (mut ws_write, mut ws_read) = socket.split();

    let mut own_session: Option<SessionHandle> = None;
    let mut subscribed_rx: Option<broadcast::Receiver<Value>> = None;
    let mut subscribed_session_id: Option<String> = None;
    let mut change_rx: Option<broadcast::Receiver<Value>> = None;

    debug!("new websocket connection");

    loop {
        tokio::select! {
            msg = ws_read.next() => {
                let text = match msg {
                    Some(Ok(Message::Text(t))) => t.to_string(),
                    Some(Ok(Message::Close(_))) | None => {
                        debug!("websocket closed by client");
                        break;
                    }
                    _ => continue,
                };

                let response = match serde_json::from_str::<Value>(&text) {
                    Ok(request) => {
                        let cmd = request.get("command").and_then(|c| c.as_str()).unwrap_or("?");
                        debug!(command = cmd, "recv command");
                        process_command(
                            &session_mgr,
                            &mut own_session,
                            &mut subscribed_rx,
                            &mut subscribed_session_id,
                            &mut change_rx,
                            &request,
                        )
                        .await
                    }
                    Err(e) => json!({
                        "command": "unknown",
                        "success": false,
                        "message": e.to_string(),
                    }),
                };

                let resp_cmd = response.get("command").and_then(|c| c.as_str()).unwrap_or("?");
                let resp_ok = response.get("success").and_then(|v| v.as_bool()).unwrap_or(false);
                debug!(command = resp_cmd, success = resp_ok, "send response");

                if ws_write
                    .send(Message::Text(response.to_string().into()))
                    .await
                    .is_err()
                {
                    warn!("ws_write failed sending response");
                    break;
                }
            }

            event = async {
                match &mut subscribed_rx {
                    Some(rx) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                match event {
                    Ok(value) => {
                        let cmd = value.get("command").and_then(|c| c.as_str()).unwrap_or("?");
                        debug!(command = cmd, "forward broadcast to subscriber");
                        if ws_write
                            .send(Message::Text(value.to_string().into()))
                            .await
                            .is_err()
                        {
                            warn!("ws_write failed sending broadcast");
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Lagged(n)) => {
                        warn!(skipped = n, "broadcast receiver lagged");
                    }
                    Err(broadcast::error::RecvError::Closed) => {
                        debug!("broadcast channel closed, clearing subscription");
                        subscribed_rx = None;
                    }
                }
            }

            event = async {
                match &mut change_rx {
                    Some(rx) => rx.recv().await,
                    None => std::future::pending().await,
                }
            } => {
                match event {
                    Ok(value) => {
                        let cmd = value.get("command").and_then(|c| c.as_str()).unwrap_or("?");
                        debug!(command = cmd, "forward change event");
                        if ws_write
                            .send(Message::Text(value.to_string().into()))
                            .await
                            .is_err()
                        {
                            warn!("ws_write failed sending change event");
                            break;
                        }
                    }
                    Err(_) => {}
                }
            }
        }
    }

    debug!("websocket handler cleaning up");

    if let Some(sid) = subscribed_session_id {
        debug!(session_id = %sid, "resetting control for watched session");
        if let Some(session) = session_mgr.get_session(&sid) {
            session.control.reset();
        }
    }

    if let Some(session) = own_session {
        debug!(session_id = %session.id, "removing owned session");
        let _ = session.broadcast_tx.send(json!({
            "command": "sessionEnded",
            "sessionId": session.id,
        }));
        session_mgr.remove_session(&session.id);
    }
}

async fn send_command(
    tx: &std::sync::mpsc::Sender<Command>,
    build: impl FnOnce(oneshot::Sender<Result<Value, String>>) -> Command,
) -> Result<Value, String> {
    let (reply_tx, reply_rx) = oneshot::channel();
    tx.send(build(reply_tx))
        .map_err(|_| "Machine thread unavailable".to_string())?;
    match reply_rx.await {
        Ok(result) => result,
        Err(_) => Err("Machine thread dropped".to_string()),
    }
}

fn get_target_session(
    session_mgr: &SessionManager,
    request: &Value,
) -> Result<SessionHandle, String> {
    let session_id = request
        .get("sessionId")
        .and_then(|v| v.as_str())
        .ok_or("Missing 'sessionId' field")?;
    session_mgr
        .get_session(session_id)
        .ok_or_else(|| format!("Session '{}' not found", session_id))
}

async fn process_command(
    session_mgr: &SessionManager,
    own_session: &mut Option<SessionHandle>,
    subscribed_rx: &mut Option<broadcast::Receiver<Value>>,
    subscribed_session_id: &mut Option<String>,
    change_rx: &mut Option<broadcast::Receiver<Value>>,
    request: &Value,
) -> Value {
    let command = request
        .get("command")
        .and_then(|c| c.as_str())
        .unwrap_or("")
        .to_uppercase();

    let result = match command.as_str() {
        "MODE" => Ok(json!({
            "command": "mode",
            "mode": "EDITOR",
            "success": true,
        })),
        "CHECK" => cmd_check(request),
        "CONVERTGRAPHML" => cmd_convert_graphml(request),
        "START" => cmd_start(session_mgr, own_session, request).await,
        "GETNEXT" => cmd_get_next(own_session).await,
        "HASNEXT" => cmd_has_next(own_session).await,
        "GETDATA" => cmd_get_data(own_session).await,
        "SETDATA" => cmd_set_data(own_session, request).await,
        "GETMODEL" => cmd_get_model(own_session).await,
        "UPDATEALLELEMENTS" => cmd_update_all_elements(own_session).await,
        "LISTSESSIONS" => cmd_list_sessions(session_mgr, change_rx),
        "SUBSCRIBESESSION" => {
            cmd_subscribe_session(session_mgr, subscribed_rx, subscribed_session_id, request).await
        }
        "UNSUBSCRIBESESSION" => {
            cmd_unsubscribe_session(session_mgr, subscribed_rx, subscribed_session_id)
        }
        "PAUSESESSION" => cmd_pause_session(session_mgr, request),
        "RESUMESESSION" => cmd_resume_session(session_mgr, request),
        "STEPSESSION" => cmd_step_session(session_mgr, request),
        "SETDELAY" => cmd_set_delay(session_mgr, request),
        "SETBREAKPOINTS" => cmd_set_breakpoints(session_mgr, request),
        _ => Ok(json!({
            "command": command.to_lowercase(),
            "success": false,
            "message": format!("Unknown command: {}", command),
        })),
    };

    match result {
        Ok(val) => val,
        Err(msg) => {
            warn!(command = %command, error = %msg, "command failed");
            json!({
                "command": command.to_lowercase(),
                "success": false,
                "message": msg,
            })
        }
    }
}

fn cmd_check(request: &Value) -> Result<Value, String> {
    let gw = request.get("gw").ok_or("Missing 'gw' field")?;
    let val = actor::handle_check(&gw.to_string())?;
    Ok(json!({
        "command": "check",
        "issues": val.get("issues").unwrap_or(&json!([])),
        "success": true,
    }))
}

fn cmd_convert_graphml(request: &Value) -> Result<Value, String> {
    let graphml = request
        .get("graphml")
        .and_then(|g| g.as_str())
        .ok_or("Missing 'graphml' field")?;
    let val = actor::handle_convert_graphml(graphml)?;
    Ok(json!({
        "command": "convertGraphml",
        "models": val.get("models").unwrap_or(&json!("")),
        "success": true,
    }))
}

async fn cmd_start(
    session_mgr: &SessionManager,
    own_session: &mut Option<SessionHandle>,
    request: &Value,
) -> Result<Value, String> {
    if let Some(prev) = own_session.take() {
        debug!(session_id = %prev.id, "removing previous session");
        session_mgr.remove_session(&prev.id);
    }

    let gw = request.get("gw").ok_or("Missing 'gw' field")?;
    let json_body = gw.to_string();
    let seed = request.get("seed").and_then(|v| v.as_u64());
    let global_data = request
        .get("globalData")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string());
    let session_name = request
        .get("name")
        .and_then(|v| v.as_str())
        .map(|s| s.to_string())
        .unwrap_or_else(|| {
            gw.get("models")
                .and_then(|m| m.as_array())
                .and_then(|arr| arr.first())
                .and_then(|m| m.get("name"))
                .and_then(|n| n.as_str())
                .unwrap_or("Unnamed")
                .to_string()
        });

    let machine_tx = spawn_machine_thread();

    let val = send_command(&machine_tx, |reply| Command::Load {
        json_body: json_body.clone(),
        seed,
        global_data,
        reply,
    })
    .await?;

    let actual_seed = val.get("seed").and_then(|v| v.as_u64());
    let handle =
        session_mgr.create_session(session_name.clone(), json_body, actual_seed, machine_tx);
    let session_id = handle.id.clone();
    *own_session = Some(handle);

    debug!(session_id = %session_id, name = %session_name, seed = ?actual_seed, "session started");

    Ok(json!({
        "command": "start",
        "success": true,
        "seed": val.get("seed").unwrap_or(&json!(0)),
        "sessionId": session_id,
    }))
}

async fn cmd_get_next(own_session: &Option<SessionHandle>) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;

    debug!(session_id = %session.id, paused = session.control.is_paused(), "getNext: entering gate");
    session.control.gate().await;
    debug!(session_id = %session.id, "getNext: gate passed");

    let val = send_command(&session.machine_tx, |reply| Command::GetNext {
        verbose: true,
        reply,
    })
    .await?;

    let model_id = val.get("modelId").and_then(|v| v.as_str()).unwrap_or("");
    let element_id = val
        .get("currentElementID")
        .and_then(|v| v.as_str())
        .unwrap_or("");
    let name = val
        .get("currentElementName")
        .and_then(|v| v.as_str())
        .unwrap_or("");

    debug!(session_id = %session.id, element = %name, "getNext: got element");

    let response = json!({
        "command": "visitedElement",
        "modelId": val.get("modelId").unwrap_or(&json!("")),
        "elementId": val.get("currentElementID").unwrap_or(&json!("")),
        "name": val.get("currentElementName").unwrap_or(&json!("")),
        "visitedCount": val.get("visitedCount").unwrap_or(&json!(0)),
        "totalCount": val.get("totalCount").unwrap_or(&json!(0)),
        "stopConditionFulfillment": val.get("stopConditionFulfillment").unwrap_or(&json!(0.0)),
        "data": val.get("data").unwrap_or(&json!("")),
        "success": true,
    });

    let subscribers = session.broadcast_tx.send(response.clone()).unwrap_or(0);
    debug!(session_id = %session.id, subscribers, "getNext: broadcast sent");

    if session
        .control
        .check_and_pause_if_breakpoint(model_id, element_id)
    {
        debug!(session_id = %session.id, element = %name, "breakpoint hit, pausing");
        let _ = session.broadcast_tx.send(json!({
            "command": "sessionPaused",
            "reason": "breakpoint",
            "modelId": model_id,
            "elementId": element_id,
        }));
    }

    let delay_ms = session.control.delay_ms();
    if delay_ms > 0 {
        debug!(session_id = %session.id, delay_ms, "getNext: sleeping");
        tokio::time::sleep(Duration::from_millis(delay_ms)).await;
    }

    Ok(response)
}

async fn cmd_has_next(own_session: &Option<SessionHandle>) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;
    let val = send_command(&session.machine_tx, |reply| Command::HasNext { reply }).await?;
    let has_next = val
        .get("hasNext")
        .and_then(|v| v.as_str())
        .unwrap_or("false");

    Ok(json!({
        "command": "hasNext",
        "hasNext": has_next == "true",
        "success": true,
    }))
}

async fn cmd_get_data(own_session: &Option<SessionHandle>) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;
    let val = send_command(&session.machine_tx, |reply| Command::GetData { reply }).await?;
    Ok(json!({
        "command": "getData",
        "data": val.get("data").unwrap_or(&json!("")),
        "success": true,
    }))
}

async fn cmd_set_data(
    own_session: &Option<SessionHandle>,
    request: &Value,
) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;
    let script = request
        .get("action")
        .and_then(|a| a.as_str())
        .ok_or("Missing 'action' field")?
        .to_string();
    send_command(&session.machine_tx, |reply| Command::SetData {
        script,
        reply,
    })
    .await?;
    Ok(json!({"command": "setData", "success": true}))
}

async fn cmd_get_model(own_session: &Option<SessionHandle>) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;
    let val = send_command(&session.machine_tx, |reply| Command::GetModel { reply }).await?;
    Ok(json!({
        "command": "getModel",
        "models": val.get("models").unwrap_or(&json!("")),
        "success": true,
    }))
}

async fn cmd_update_all_elements(own_session: &Option<SessionHandle>) -> Result<Value, String> {
    let session = own_session.as_ref().ok_or("No active session")?;
    let val = send_command(&session.machine_tx, |reply| Command::UpdateAllElements {
        reply,
    })
    .await?;
    Ok(json!({
        "command": "updateAllElements",
        "elements": val.get("elements").unwrap_or(&json!([])),
        "success": true,
    }))
}

fn cmd_list_sessions(
    session_mgr: &SessionManager,
    change_rx: &mut Option<broadcast::Receiver<Value>>,
) -> Result<Value, String> {
    if change_rx.is_none() {
        *change_rx = Some(session_mgr.subscribe_changes());
    }

    let sessions = session_mgr.list_sessions();
    let list: Vec<Value> = sessions
        .iter()
        .map(|s| {
            json!({
                "id": s.id,
                "name": s.name,
            })
        })
        .collect();

    debug!(count = sessions.len(), "listSessions");

    Ok(json!({
        "command": "sessions",
        "sessions": list,
        "success": true,
    }))
}

async fn cmd_subscribe_session(
    session_mgr: &SessionManager,
    subscribed_rx: &mut Option<broadcast::Receiver<Value>>,
    subscribed_session_id: &mut Option<String>,
    request: &Value,
) -> Result<Value, String> {
    if let Some(prev_id) = subscribed_session_id.take() {
        debug!(prev_session = %prev_id, "resetting previous subscription");
        if let Some(prev_session) = session_mgr.get_session(&prev_id) {
            prev_session.control.reset();
        }
    }

    let session_id = request
        .get("sessionId")
        .and_then(|v| v.as_str())
        .ok_or("Missing 'sessionId' field")?;

    let session = session_mgr
        .get_session(session_id)
        .ok_or_else(|| format!("Session '{}' not found", session_id))?;

    debug!(session_id = %session_id, "fetching element snapshot");

    let elements = send_command(&session.machine_tx, |reply| Command::UpdateAllElements {
        reply,
    })
    .await
    .ok()
    .and_then(|v| v.get("elements").cloned())
    .unwrap_or(json!([]));

    let models_value = send_command(&session.machine_tx, |reply| Command::GetModel { reply })
        .await
        .ok()
        .and_then(|v| v.get("models").and_then(|m| m.as_str()).map(String::from))
        .and_then(|s| serde_json::from_str::<Value>(&s).ok())
        .unwrap_or_else(|| {
            serde_json::from_str::<Value>(&session.model_json).unwrap_or(json!(null))
        });

    let elem_count = elements.as_array().map(|a| a.len()).unwrap_or(0);
    debug!(session_id = %session_id, elements = elem_count, "subscribing to broadcast");

    *subscribed_rx = Some(session.broadcast_tx.subscribe());
    *subscribed_session_id = Some(session_id.to_string());

    Ok(json!({
        "command": "subscribeSession",
        "sessionId": session.id,
        "name": session.name,
        "models": models_value,
        "elements": elements,
        "seed": session.seed,
        "paused": session.control.is_paused(),
        "success": true,
    }))
}

fn cmd_unsubscribe_session(
    session_mgr: &SessionManager,
    subscribed_rx: &mut Option<broadcast::Receiver<Value>>,
    subscribed_session_id: &mut Option<String>,
) -> Result<Value, String> {
    if let Some(sid) = subscribed_session_id.take() {
        debug!(session_id = %sid, "unsubscribing, resetting control");
        if let Some(session) = session_mgr.get_session(&sid) {
            session.control.reset();
        }
    }
    *subscribed_rx = None;
    Ok(json!({
        "command": "unsubscribeSession",
        "success": true,
    }))
}

fn cmd_pause_session(session_mgr: &SessionManager, request: &Value) -> Result<Value, String> {
    let session = get_target_session(session_mgr, request)?;
    debug!(session_id = %session.id, "pauseSession");
    session.control.pause();
    let _ = session.broadcast_tx.send(json!({
        "command": "sessionPaused",
        "sessionId": session.id,
    }));
    Ok(json!({"command": "pauseSession", "success": true}))
}

fn cmd_resume_session(session_mgr: &SessionManager, request: &Value) -> Result<Value, String> {
    let session = get_target_session(session_mgr, request)?;
    debug!(session_id = %session.id, "resumeSession");
    session.control.resume();
    let _ = session.broadcast_tx.send(json!({
        "command": "sessionResumed",
        "sessionId": session.id,
    }));
    Ok(json!({"command": "resumeSession", "success": true}))
}

fn cmd_step_session(session_mgr: &SessionManager, request: &Value) -> Result<Value, String> {
    let session = get_target_session(session_mgr, request)?;
    debug!(session_id = %session.id, "stepSession (pause + step)");
    session.control.pause();
    session.control.step();
    Ok(json!({"command": "stepSession", "success": true}))
}

fn cmd_set_delay(session_mgr: &SessionManager, request: &Value) -> Result<Value, String> {
    let session = get_target_session(session_mgr, request)?;
    let ms = request.get("value").and_then(|v| v.as_u64()).unwrap_or(0);
    debug!(session_id = %session.id, delay_ms = ms, "setDelay");
    session.control.set_delay(ms);
    Ok(json!({"command": "setDelay", "success": true}))
}

fn cmd_set_breakpoints(session_mgr: &SessionManager, request: &Value) -> Result<Value, String> {
    let session = get_target_session(session_mgr, request)?;
    let bps: HashSet<String> = request
        .get("breakpoints")
        .and_then(|v| v.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect()
        })
        .unwrap_or_default();
    debug!(session_id = %session.id, count = bps.len(), "setBreakpoints");
    session.control.set_breakpoints(bps);
    Ok(json!({"command": "setBreakpoints", "success": true}))
}