PicoBot/src/channels/cli_chat.rs

use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::{mpsc, Mutex};
use uuid::Uuid;

use crate::bus::{ControlInbound, ControlMessage, ControlOutbound, InboundMessage, MessageBus, OutboundMessage};
use crate::protocol::{parse_inbound, WsInbound, WsOutbound};

use super::base::{Channel, ChannelError};

// ============================================================================
// Client - Connected CLI client
// ============================================================================

pub(crate) struct Client {
    sender: mpsc::Sender<WsOutbound>,
    current_session_id: Mutex<Option<String>>,
}

// ============================================================================
// CliChatChannel - Channel implementation for CLI chat
// ============================================================================

pub struct CliChatChannel {
    bus: std::sync::Mutex<Option<Arc<MessageBus>>>,
    clients: Mutex<Vec<Arc<Client>>>,
}

impl CliChatChannel {
    pub fn new() -> Self {
        Self {
            bus: std::sync::Mutex::new(None),
            clients: Mutex::new(Vec::new()),
        }
    }

    /// Register a new client connection, returns (session_id, client)
    pub(crate) async fn register_client(&self, sender: mpsc::Sender<WsOutbound>) -> (String, Arc<Client>) {
        let client = Arc::new(Client {
            sender,
            current_session_id: Mutex::new(None),
        });
        self.clients.lock().await.push(client.clone());

        // Create initial session via control message
        let session_id = match self.create_session_via_control(None).await {
            Ok(id) => id,
            Err(e) => {
                tracing::error!(error = %e, "Failed to create initial session");
                Uuid::new_v4().to_string()
            }
        };

        // Set current session id in client
        {
            let mut current = client.current_session_id.lock().await;
            *current = Some(session_id.clone());
        }

        (session_id, client)
    }

    /// Handle an inbound message from a client
    pub(crate) async fn handle_inbound(&self, client: Arc<Client>, raw_msg: &str) {
        match parse_inbound(raw_msg) {
            Ok(inbound) => {
                match self.handle_ws_inbound(client.clone(), inbound).await {
                    Ok(()) => {}
                    Err(e) => {
                        tracing::warn!(error = %e, "Failed to handle inbound message");
                        let _ = client
                            .sender
                            .send(WsOutbound::Error {
                                code: "INTERNAL_ERROR".to_string(),
                                message: e.to_string(),
                            })
                            .await;
                    }
                }
            }
            Err(e) => {
                tracing::warn!(error = %e, "Failed to parse inbound message");
                let _ = client
                    .sender
                    .send(WsOutbound::Error {
                        code: "PARSE_ERROR".to_string(),
                        message: e.to_string(),
                    })
                    .await;
            }
        }
    }

    async fn handle_ws_inbound(&self, client: Arc<Client>, inbound: WsInbound) -> Result<(), ChannelError> {
        let bus = {
            let guard = self.bus.lock().unwrap();
            guard.clone().ok_or_else(|| ChannelError::Other("Channel not started".to_string()))?
        };

        let mut current_session_guard = client.current_session_id.lock().await;

        match inbound {
            WsInbound::UserInput { content, chat_id, .. } => {
                let chat_id = chat_id.or(current_session_guard.clone()).unwrap_or_else(|| Uuid::new_v4().to_string());

                // If no session, create one first
                if current_session_guard.is_none() {
                    let new_id = self.create_session_via_control(None).await?;
                    *current_session_guard = Some(new_id);
                }

                let session_id = current_session_guard.clone().unwrap();

                // Publish to bus for AI processing
                let msg = InboundMessage {
                    channel: self.name().to_string(),
                    sender_id: "cli".to_string(),
                    chat_id: session_id.clone(),
                    content,
                    timestamp: current_timestamp(),
                    media: Vec::new(),
                    metadata: Default::default(),
                    forwarded_metadata: Default::default(),
                };
                bus.publish_inbound(msg).await?;
            }
            WsInbound::ClearHistory { chat_id, session_id } => {
                let target = session_id
                    .or(chat_id)
                    .or(current_session_guard.clone())
                    .ok_or_else(|| ChannelError::Other("No active session".to_string()))?;

                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::ClearHistory { session_id: target.clone() },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::HistoryCleared { .. })) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::HistoryCleared { session_id: target })
                            .await;
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type, ignore
                    }
                    Some(Err(e)) => {
                        return Err(e);
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::CreateSession { title } => {
                let new_id = self.create_session_via_control(title.as_deref()).await?;
                *current_session_guard = Some(new_id.clone());
                let _ = client
                    .sender
                    .send(WsOutbound::SessionCreated {
                        session_id: new_id,
                        title: title.unwrap_or_default(),
                    })
                    .await;
            }
            WsInbound::ListSessions { include_archived } => {
                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::ListSessions { include_archived },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::SessionList { sessions })) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::SessionList {
                                sessions,
                                current_session_id: current_session_guard.clone(),
                            })
                            .await;
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type
                    }
                    Some(Err(e)) => {
                        return Err(e);
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::LoadSession { session_id } => {
                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::LoadSession { session_id: session_id.clone() },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::SessionLoaded { session_id, title, message_count })) => {
                        *current_session_guard = Some(session_id.clone());
                        let _ = client
                            .sender
                            .send(WsOutbound::SessionLoaded {
                                session_id,
                                title,
                                message_count,
                            })
                            .await;
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type
                    }
                    Some(Err(e)) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::Error {
                                code: "SESSION_NOT_FOUND".to_string(),
                                message: format!("Session not found: {}", session_id),
                            })
                            .await;
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::RenameSession { session_id, title } => {
                let target = session_id.or(current_session_guard.clone()).ok_or_else(|| {
                    ChannelError::Other("No active session".to_string())
                })?;

                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::RenameSession { session_id: target.clone(), title: title.clone() },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::SessionRenamed { session_id, title })) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::SessionRenamed { session_id, title })
                            .await;
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type
                    }
                    Some(Err(e)) => {
                        return Err(e);
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::ArchiveSession { session_id } => {
                let target = session_id.or(current_session_guard.clone()).ok_or_else(|| {
                    ChannelError::Other("No active session".to_string())
                })?;

                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::ArchiveSession { session_id: target.clone() },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::SessionArchived { session_id })) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::SessionArchived { session_id })
                            .await;
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type
                    }
                    Some(Err(e)) => {
                        return Err(e);
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::DeleteSession { session_id } => {
                let target = session_id.or(current_session_guard.clone()).ok_or_else(|| {
                    ChannelError::Other("No active session".to_string())
                })?;

                let (reply_tx, mut reply_rx) = mpsc::channel(1);
                bus.publish_control(ControlMessage {
                    op: ControlInbound::DeleteSession { session_id: target.clone() },
                    reply_tx,
                }).await?;

                match reply_rx.recv().await {
                    Some(Ok(ControlOutbound::SessionDeleted { session_id })) => {
                        let _ = client
                            .sender
                            .send(WsOutbound::SessionDeleted { session_id: session_id.clone() })
                            .await;

                        // If deleting current session, create a new one
                        if current_session_guard.as_deref() == Some(&target) {
                            drop(reply_rx);
                            if let Ok(new_id) = self.create_session_via_control(None).await {
                                *current_session_guard = Some(new_id.clone());
                                let _ = client
                                    .sender
                                    .send(WsOutbound::SessionCreated {
                                        session_id: new_id,
                                        title: String::new(),
                                    })
                                    .await;
                            }
                        }
                    }
                    Some(Ok(_)) => {
                        // Unexpected response type
                    }
                    Some(Err(e)) => {
                        return Err(e);
                    }
                    None => {
                        return Err(ChannelError::Other("Control channel closed".to_string()));
                    }
                }
            }
            WsInbound::Ping => {
                let _ = client.sender.send(WsOutbound::Pong).await;
            }
        }
        Ok(())
    }

    /// Create a session via control message and return the session_id
    async fn create_session_via_control(&self, title: Option<&str>) -> Result<String, ChannelError> {
        let bus = {
            let guard = self.bus.lock().unwrap();
            guard.clone().ok_or_else(|| ChannelError::Other("Channel not started".to_string()))?
        };

        let (reply_tx, mut reply_rx) = mpsc::channel(1);
        bus.publish_control(ControlMessage {
            op: ControlInbound::CreateSession { title: title.map(String::from) },
            reply_tx,
        }).await?;

        match reply_rx.recv().await {
            Some(Ok(ControlOutbound::SessionCreated { session_id, .. })) => {
                Ok(session_id)
            }
            Some(Ok(_)) => {
                Err(ChannelError::Other("Unexpected response type".to_string()))
            }
            Some(Err(e)) => Err(e),
            None => Err(ChannelError::Other("Control channel closed".to_string())),
        }
    }
}

#[async_trait]
impl Channel for CliChatChannel {
    fn name(&self) -> &str {
        "cli_chat"
    }

    fn is_running(&self) -> bool {
        self.bus.lock().unwrap().is_some()
    }

    async fn start(&self, bus: Arc<MessageBus>) -> Result<(), ChannelError> {
        *self.bus.lock().unwrap() = Some(bus);
        Ok(())
    }

    async fn stop(&self) -> Result<(), ChannelError> {
        *self.bus.lock().unwrap() = None;
        Ok(())
    }

    async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
        let clients = self.clients.lock().await.clone();
        for client in clients {
            let outbound = WsOutbound::AssistantResponse {
                id: Uuid::new_v4().to_string(),
                content: msg.content.clone(),
                role: "assistant".to_string(),
            };
            let _ = client.sender.send(outbound).await;
        }
        Ok(())
    }
}

fn current_timestamp() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64
}