path: root/makima/src/daemon/ws/client.rs

//! WebSocket client for connecting to the makima server.

use std::sync::Arc;
use std::time::Duration;

use backoff::backoff::Backoff;
use backoff::ExponentialBackoff;
use futures::{SinkExt, StreamExt};
use tokio::sync::{mpsc, RwLock};
use tokio_tungstenite::{connect_async, tungstenite::{client::IntoClientRequest, Message}};
use uuid::Uuid;

use super::protocol::{DaemonCommand, DaemonMessage};
use crate::daemon::config::ServerConfig;
use crate::daemon::error::{DaemonError, Result};

/// WebSocket client state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Not connected to server.
    Disconnected,
    /// Currently connecting.
    Connecting,
    /// Connected and authenticated.
    Connected,
    /// Connection failed, will retry.
    Reconnecting,
    /// Permanently failed (e.g., auth failure).
    Failed,
}

/// WebSocket client for daemon-server communication.
pub struct WsClient {
    config: ServerConfig,
    machine_id: String,
    hostname: String,
    max_concurrent_tasks: i32,
    state: Arc<RwLock<ConnectionState>>,
    daemon_id: Arc<RwLock<Option<Uuid>>>,
    /// Channel to receive messages to send to server.
    outgoing_rx: mpsc::Receiver<DaemonMessage>,
    /// Sender for outgoing messages (clone this to send messages).
    outgoing_tx: mpsc::Sender<DaemonMessage>,
    /// Channel to send received commands to the task manager.
    incoming_tx: mpsc::Sender<DaemonCommand>,
}

impl WsClient {
    /// Create a new WebSocket client.
    pub fn new(
        config: ServerConfig,
        machine_id: String,
        hostname: String,
        max_concurrent_tasks: i32,
        incoming_tx: mpsc::Sender<DaemonCommand>,
    ) -> Self {
        let (outgoing_tx, outgoing_rx) = mpsc::channel(256);

        Self {
            config,
            machine_id,
            hostname,
            max_concurrent_tasks,
            state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
            daemon_id: Arc::new(RwLock::new(None)),
            outgoing_rx,
            outgoing_tx,
            incoming_tx,
        }
    }

    /// Get a sender for outgoing messages.
    pub fn sender(&self) -> mpsc::Sender<DaemonMessage> {
        self.outgoing_tx.clone()
    }

    /// Get current connection state.
    pub async fn state(&self) -> ConnectionState {
        *self.state.read().await
    }

    /// Get daemon ID if authenticated.
    pub async fn daemon_id(&self) -> Option<Uuid> {
        *self.daemon_id.read().await
    }

    /// Run the WebSocket client with automatic reconnection.
    pub async fn run(&mut self) -> Result<()> {
        let mut backoff = ExponentialBackoff {
            initial_interval: Duration::from_secs(self.config.reconnect_interval_secs),
            max_interval: Duration::from_secs(60),
            max_elapsed_time: if self.config.max_reconnect_attempts > 0 {
                Some(Duration::from_secs(
                    self.config.reconnect_interval_secs * self.config.max_reconnect_attempts as u64 * 10,
                ))
            } else {
                None // Infinite retries
            },
            ..Default::default()
        };

        loop {
            *self.state.write().await = ConnectionState::Connecting;
            tracing::info!("Connecting to server: {}", self.config.url);

            match self.connect_and_run().await {
                Ok(()) => {
                    // Clean shutdown
                    tracing::info!("WebSocket connection closed cleanly");
                    break;
                }
                Err(DaemonError::AuthFailed(msg)) => {
                    tracing::error!("Authentication failed: {}", msg);
                    *self.state.write().await = ConnectionState::Failed;
                    return Err(DaemonError::AuthFailed(msg));
                }
                Err(e) => {
                    tracing::warn!("Connection error: {}", e);
                    *self.state.write().await = ConnectionState::Reconnecting;

                    if let Some(delay) = backoff.next_backoff() {
                        tracing::info!("Reconnecting in {:?}...", delay);
                        tokio::time::sleep(delay).await;
                    } else {
                        tracing::error!("Max reconnection attempts reached");
                        *self.state.write().await = ConnectionState::Failed;
                        return Err(DaemonError::ConnectionLost);
                    }
                }
            }
        }

        Ok(())
    }

    /// Connect to server and run the message loop.
    async fn connect_and_run(&mut self) -> Result<()> {
        // Build WebSocket URL
        let ws_url = format!("{}/api/v1/mesh/daemons/connect", self.config.url);
        tracing::debug!("Connecting to WebSocket: {}", ws_url);

        // Build request with API key header
        let mut request = ws_url.into_client_request()?;
        request.headers_mut().insert(
            "x-makima-api-key",
            self.config.api_key.parse().map_err(|_| {
                DaemonError::AuthFailed("Invalid API key format".into())
            })?,
        );

        // Connect with API key in headers
        let (ws_stream, _response) = connect_async(request).await?;
        let (mut write, mut read) = ws_stream.split();

        // Send daemon info after connection (server authenticated us via header)
        let info_msg = DaemonMessage::authenticate(
            &self.config.api_key,
            &self.machine_id,
            &self.hostname,
            self.max_concurrent_tasks,
        );
        let info_json = serde_json::to_string(&info_msg)?;
        write.send(Message::Text(info_json)).await?;

        // Wait for authentication response
        let auth_response = read
            .next()
            .await
            .ok_or(DaemonError::ConnectionLost)??;

        let auth_text = match auth_response {
            Message::Text(text) => text,
            Message::Close(_) => return Err(DaemonError::ConnectionLost),
            _ => return Err(DaemonError::AuthFailed("Unexpected response type".into())),
        };

        let command: DaemonCommand = serde_json::from_str(&auth_text)?;
        match command {
            DaemonCommand::Authenticated { daemon_id } => {
                tracing::info!("Authenticated with daemon ID: {}", daemon_id);
                *self.daemon_id.write().await = Some(daemon_id);
                *self.state.write().await = ConnectionState::Connected;

                // Send daemon directories info to server
                let working_directory = std::env::current_dir()
                    .map(|p| p.to_string_lossy().to_string())
                    .unwrap_or_else(|_| ".".to_string());
                let home_directory = dirs::home_dir()
                    .map(|h| h.join(".makima").join("home"))
                    .unwrap_or_else(|| std::path::PathBuf::from("~/.makima/home"));
                // Create home directory if it doesn't exist
                if let Err(e) = std::fs::create_dir_all(&home_directory) {
                    tracing::warn!("Failed to create home directory {:?}: {}", home_directory, e);
                }
                let home_directory_str = home_directory.to_string_lossy().to_string();
                let worktrees_directory = dirs::home_dir()
                    .map(|h| h.join(".makima").join("worktrees").to_string_lossy().to_string())
                    .unwrap_or_else(|| "~/.makima/worktrees".to_string());

                let dirs_msg = DaemonMessage::DaemonDirectories {
                    working_directory,
                    home_directory: home_directory_str,
                    worktrees_directory,
                };
                let dirs_json = serde_json::to_string(&dirs_msg)?;
                write.send(Message::Text(dirs_json)).await?;
                tracing::info!("Sent daemon directories info to server");
            }
            DaemonCommand::Error { code, message } => {
                return Err(DaemonError::AuthFailed(format!("{}: {}", code, message)));
            }
            _ => {
                return Err(DaemonError::AuthFailed(
                    "Unexpected response to authentication".into(),
                ));
            }
        }

        // Start main message loop
        let heartbeat_interval = Duration::from_secs(self.config.heartbeat_interval_secs);
        let mut heartbeat_timer = tokio::time::interval(heartbeat_interval);

        loop {
            tokio::select! {
                // Handle incoming server commands
                msg = read.next() => {
                    match msg {
                        Some(Ok(Message::Text(text))) => {
                            tracing::info!("Received WebSocket message: {} bytes", text.len());
                            match serde_json::from_str::<DaemonCommand>(&text) {
                                Ok(command) => {
                                    tracing::info!("Parsed command: {:?}", command);
                                    tracing::info!("Sending command to task manager channel...");
                                    if self.incoming_tx.send(command).await.is_err() {
                                        tracing::warn!("Command receiver dropped, shutting down");
                                        break;
                                    }
                                    tracing::info!("Command sent to task manager successfully");
                                }
                                Err(e) => {
                                    tracing::warn!("Failed to parse server message: {}", e);
                                    tracing::debug!("Raw message: {}", text);
                                }
                            }
                        }
                        Some(Ok(Message::Ping(data))) => {
                            write.send(Message::Pong(data)).await?;
                        }
                        Some(Ok(Message::Close(_))) | None => {
                            tracing::info!("Server closed connection");
                            return Err(DaemonError::ConnectionLost);
                        }
                        Some(Err(e)) => {
                            tracing::warn!("WebSocket error: {}", e);
                            return Err(e.into());
                        }
                        _ => {}
                    }
                }

                // Handle outgoing messages
                msg = self.outgoing_rx.recv() => {
                    match msg {
                        Some(message) => {
                            let json = serde_json::to_string(&message)?;
                            tracing::trace!("Sending message: {}", json);
                            write.send(Message::Text(json)).await?;
                        }
                        None => {
                            // Sender dropped, shutdown
                            tracing::info!("Outgoing channel closed, shutting down");
                            break;
                        }
                    }
                }

                // Send heartbeat
                _ = heartbeat_timer.tick() => {
                    // Get active task IDs from task manager
                    // For now, send empty list - will be connected to task manager
                    let heartbeat = DaemonMessage::heartbeat(vec![]);
                    let json = serde_json::to_string(&heartbeat)?;
                    write.send(Message::Text(json)).await?;
                }
            }
        }

        Ok(())
    }
}