PicoBot/src/client/mod.rs

pub use crate::protocol::{WsInbound, WsOutbound, serialize_inbound, serialize_outbound};

use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::{connect_async, tungstenite::Message};

use crate::cli::{InputCommand, InputEvent, InputHandler};

fn format_session_list(
    sessions: &[crate::protocol::SessionSummary],
    current_session_id: Option<&str>,
) -> String {
    if sessions.is_empty() {
        return "No sessions found.".to_string();
    }

    let mut lines = Vec::with_capacity(sessions.len() + 1);
    lines.push("Sessions:".to_string());
    for session in sessions {
        let marker = if current_session_id == Some(session.session_id.as_str()) {
            "*"
        } else {
            "-"
        };
        let archived = if session.archived_at.is_some() {
            " [archived]"
        } else {
            ""
        };
        lines.push(format!(
            "{} {} | {} | {} messages{}",
            marker, session.session_id, session.title, session.message_count, archived,
        ));
    }

    lines.join("\n")
}

fn parse_message(raw: &str) -> Result<WsOutbound, serde_json::Error> {
    serde_json::from_str(raw)
}

fn format_json(value: &serde_json::Value) -> String {
    serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string())
}

pub async fn run(gateway_url: &str) -> Result<(), Box<dyn std::error::Error>> {
    let (ws_stream, _) = connect_async(gateway_url).await?;
    tracing::info!(url = %gateway_url, "Connected to gateway");

    let (mut sender, mut receiver) = ws_stream.split();

    let mut input = InputHandler::new();
    let mut current_session_id: Option<String> = None;
    input.write_output("picobot CLI - Commands: /new [title], /reset, /sessions, /use <session>, /rename <title>, /archive, /delete, /clear, /quit\n").await?;

    // Main loop: poll both stdin and WebSocket
    loop {
        tokio::select! {
            // Handle WebSocket messages
            msg = receiver.next() => {
                match msg {
                    Some(Ok(Message::Text(text))) => {
                        let text = text.to_string();
                        if let Ok(outbound) = parse_message(&text) {
                            match outbound {
                                WsOutbound::AssistantResponse { content, .. } => {
                                    input.write_response(&content).await?;
                                }
                                WsOutbound::ToolCall { tool_name, arguments, .. } => {
                                    input.write_output(&format!("Tool call: {}\n{}\n", tool_name, format_json(&arguments))).await?;
                                }
                                WsOutbound::ToolResult { tool_name, content, .. } => {
                                    input.write_output(&format!("Tool result: {}\n{}\n", tool_name, content)).await?;
                                }
                                WsOutbound::ToolPending { tool_name, content, resume_hint, .. } => {
                                    input.write_output(&format!("Tool pending: {}\n{}\n{}\n", tool_name, content, resume_hint)).await?;
                                }
                                WsOutbound::Error { message, .. } => {
                                    input.write_output(&format!("Error: {}", message)).await?;
                                }
                                WsOutbound::SessionEstablished { session_id } => {
                                    current_session_id = Some(session_id.clone());
                                    #[cfg(debug_assertions)]
                                    tracing::debug!(session_id = %session_id, "Session established");
                                    input.write_output(&format!("Session: {}\n", session_id)).await?;
                                }
                                WsOutbound::SessionCreated { session_id, title } => {
                                    current_session_id = Some(session_id.clone());
                                    input.write_output(&format!("Created session: {} ({})\n", session_id, title)).await?;
                                }
                                WsOutbound::SessionList { sessions, current_session_id: listed_current } => {
                                    let display = format_session_list(&sessions, listed_current.as_deref());
                                    input.write_output(&format!("{}\n", display)).await?;
                                }
                                WsOutbound::SessionLoaded { session_id, title, message_count } => {
                                    current_session_id = Some(session_id.clone());
                                    input.write_output(&format!("Loaded session: {} ({}, {} messages)\n", session_id, title, message_count)).await?;
                                }
                                WsOutbound::SessionRenamed { session_id, title } => {
                                    input.write_output(&format!("Renamed session: {} -> {}\n", session_id, title)).await?;
                                }
                                WsOutbound::SessionArchived { session_id } => {
                                    input.write_output(&format!("Archived session: {}\n", session_id)).await?;
                                }
                                WsOutbound::SessionDeleted { session_id } => {
                                    if current_session_id.as_deref() == Some(session_id.as_str()) {
                                        current_session_id = None;
                                    }
                                    input.write_output(&format!("Deleted session: {}\n", session_id)).await?;
                                }
                                WsOutbound::HistoryCleared { session_id } => {
                                    input.write_output(&format!("Cleared history for session: {}\n", session_id)).await?;
                                }
                                _ => {}
                            }
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => {
                        tracing::info!("Gateway disconnected");
                        input.write_output("Gateway disconnected").await?;
                        break;
                    }
                    _ => {}
                }
            }
            // Handle stdin input
            result = input.read_input("> ") => {
                match result {
                    Ok(Some(event)) => {
                        match event {
                            InputEvent::Command(InputCommand::Exit) => {
                                input.write_output("Goodbye!").await?;
                                break;
                            }
                            InputEvent::Command(InputCommand::Clear) => {
                                let inbound = WsInbound::ClearHistory {
                                    chat_id: None,
                                    session_id: current_session_id.clone(),
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::New(title)) => {
                                let inbound = WsInbound::CreateSession { title };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Sessions) => {
                                let inbound = WsInbound::ListSessions {
                                    include_archived: true,
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Use(session_id)) => {
                                let inbound = WsInbound::LoadSession { session_id };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Rename(title)) => {
                                let inbound = WsInbound::RenameSession {
                                    session_id: current_session_id.clone(),
                                    title,
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Archive) => {
                                let inbound = WsInbound::ArchiveSession {
                                    session_id: current_session_id.clone(),
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Delete) => {
                                let inbound = WsInbound::DeleteSession {
                                    session_id: current_session_id.clone(),
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    let _ = sender.send(Message::Text(text.into())).await;
                                }
                                continue;
                            }
                            InputEvent::Message(msg) => {
                                let inbound = WsInbound::UserInput {
                                    content: msg.content,
                                    channel: None,
                                    chat_id: current_session_id.clone(),
                                    sender_id: None,
                                };
                                if let Ok(text) = serialize_inbound(&inbound) {
                                    if sender.send(Message::Text(text.into())).await.is_err() {
                                        tracing::error!("Failed to send message to gateway");
                                        break;
                                    }
                                }
                            }
                        }
                    }
                    Ok(None) => break,
                    Err(e) => {
                        tracing::error!(error = %e, "Input error");
                        break;
                    }
                }
            }
        }
    }

    Ok(())
}