mechutil 0.8.0

//
// Copyright (C) 2024 - 2025 Automated Design Corp. All Rights Reserved.
//

//! IPC Server for autocore-server to accept connections from external modules.
//!
//! This module provides the server-side infrastructure for accepting and managing
//! connections from external modules. It handles:
//! - Accepting new TCP connections
//! - Module registration handshake
//! - Routing messages to/from connected modules
//! - Connection lifecycle management
//!
//! # Example
//!
//! ```ignore
//! use mechutil::ipc::{IpcServer, IpcServerConfig, ModuleConnection};
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let config = IpcServerConfig::new("127.0.0.1:9100");
//!     let server = IpcServer::bind(config).await?;
//!
//!     // Accept connections in a loop
//!     while let Ok(connection) = server.accept().await {
//!         tokio::spawn(async move {
//!             handle_module_connection(connection).await;
//!         });
//!     }
//!
//!     Ok(())
//! }
//! ```

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{broadcast, mpsc, Mutex, RwLock};

use super::command_message::{CommandMessage, MessageType};
use super::error::IpcError;
use super::message::ModuleRegistration;
use super::transport::SplitTransport;

/// Configuration for the IPC server.
#[derive(Debug, Clone)]
pub struct IpcServerConfig {
    /// Address to bind to
    pub bind_addr: String,
    /// Maximum number of concurrent connections
    pub max_connections: usize,
    /// Heartbeat timeout (connection considered dead if no heartbeat received)
    pub heartbeat_timeout: Duration,
    /// Registration timeout
    pub registration_timeout: Duration,
}

impl Default for IpcServerConfig {
    fn default() -> Self {
        Self {
            bind_addr: "127.0.0.1:9100".to_string(),
            max_connections: 64,
            heartbeat_timeout: Duration::from_secs(30),
            registration_timeout: Duration::from_secs(10),
        }
    }
}

impl IpcServerConfig {
    pub fn new(bind_addr: &str) -> Self {
        Self {
            bind_addr: bind_addr.to_string(),
            ..Default::default()
        }
    }
}

/// Represents a connected module.
#[derive(Debug)]
pub struct ModuleConnection {
    /// Unique connection ID
    pub id: usize,
    /// Module registration info
    pub registration: ModuleRegistration,
    /// Transport for communication
    transport: SplitTransport,
    /// Whether the connection is active
    active: Arc<AtomicBool>,
    /// Channel for sending messages to this module
    outbound_tx: mpsc::Sender<CommandMessage>,
    /// Last heartbeat timestamp
    last_heartbeat: Arc<Mutex<std::time::Instant>>,
}

impl ModuleConnection {
    /// Create a new module connection.
    fn new(
        id: usize,
        registration: ModuleRegistration,
        transport: SplitTransport,
    ) -> (Self, mpsc::Receiver<CommandMessage>) {
        let (outbound_tx, outbound_rx) = mpsc::channel(256);

        let conn = Self {
            id,
            registration,
            transport,
            active: Arc::new(AtomicBool::new(true)),
            outbound_tx,
            last_heartbeat: Arc::new(Mutex::new(std::time::Instant::now())),
        };

        (conn, outbound_rx)
    }

    /// Get the module's domain name.
    pub fn domain(&self) -> &str {
        &self.registration.name
    }

    /// Check if the connection is active.
    pub fn is_active(&self) -> bool {
        self.active.load(Ordering::SeqCst)
    }

    /// Mark the connection as inactive.
    pub fn deactivate(&self) {
        self.active.store(false, Ordering::SeqCst);
    }

    /// Send a message to this module.
    pub async fn send(&self, msg: CommandMessage) -> Result<(), IpcError> {
        if !self.is_active() {
            return Err(IpcError::Connection("Connection is not active".to_string()));
        }

        self.outbound_tx
            .send(msg)
            .await
            .map_err(|e| IpcError::Channel(e.to_string()))
    }

    /// Receive a message from this module.
    pub async fn recv(&self) -> Result<CommandMessage, IpcError> {
        if !self.is_active() {
            return Err(IpcError::Connection("Connection is not active".to_string()));
        }

        self.transport.recv().await
    }

    /// Update the heartbeat timestamp.
    pub async fn update_heartbeat(&self) {
        let mut last = self.last_heartbeat.lock().await;
        *last = std::time::Instant::now();
    }

    /// Check if the connection has timed out.
    pub async fn is_timed_out(&self, timeout: Duration) -> bool {
        let last = self.last_heartbeat.lock().await;
        last.elapsed() > timeout
    }

    /// Clone the transport for use in separate tasks.
    pub fn clone_transport(&self) -> SplitTransport {
        self.transport.clone()
    }

    /// Get sender for outbound messages.
    pub fn outbound_sender(&self) -> mpsc::Sender<CommandMessage> {
        self.outbound_tx.clone()
    }
}

/// Connection ID counter
static CONNECTION_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// Get a unique connection ID
fn next_connection_id() -> usize {
    CONNECTION_ID_COUNTER.fetch_add(1, Ordering::SeqCst)
}

/// IPC Server that accepts connections from external modules.
pub struct IpcServer {
    config: IpcServerConfig,
    listener: TcpListener,
    connections: Arc<RwLock<HashMap<usize, Arc<ModuleConnection>>>>,
    /// Map from domain name to connection ID
    domain_map: Arc<RwLock<HashMap<String, usize>>>,
    running: Arc<AtomicBool>,
    /// Broadcast channel for server-wide events
    event_tx: broadcast::Sender<ServerEvent>,
}

/// Events emitted by the server.
#[derive(Debug, Clone)]
pub enum ServerEvent {
    /// A new module connected
    ModuleConnected {
        connection_id: usize,
        domain: String,
    },
    /// A module disconnected
    ModuleDisconnected {
        connection_id: usize,
        domain: String,
    },
    /// A message was received from a module
    MessageReceived {
        connection_id: usize,
        domain: String,
        message: CommandMessage,
    },
}

impl IpcServer {
    /// Bind to the configured address and create a new server.
    pub async fn bind(config: IpcServerConfig) -> Result<Self, IpcError> {
        let listener = TcpListener::bind(&config.bind_addr)
            .await
            .map_err(|e| IpcError::Connection(format!("Failed to bind: {}", e)))?;

        log::info!("IPC server listening on {}", config.bind_addr);

        let (event_tx, _) = broadcast::channel(256);

        Ok(Self {
            config,
            listener,
            connections: Arc::new(RwLock::new(HashMap::new())),
            domain_map: Arc::new(RwLock::new(HashMap::new())),
            running: Arc::new(AtomicBool::new(true)),
            event_tx,
        })
    }

    /// Accept a new connection from a module.
    ///
    /// This performs the registration handshake and returns the connection
    /// if successful.
    pub async fn accept(&self) -> Result<Arc<ModuleConnection>, IpcError> {
        loop {
            let (stream, addr) = self.listener.accept().await?;

            log::info!("New connection from {}", addr);

            // Check connection limit
            {
                let connections = self.connections.read().await;
                if connections.len() >= self.config.max_connections {
                    log::warn!("Connection limit reached, rejecting {}", addr);
                    continue;
                }
            }

            // Disable Nagle's algorithm
            stream.set_nodelay(true).ok();

            // Perform registration handshake
            match self.perform_registration(stream).await {
                Ok(connection) => {
                    let connection = Arc::new(connection);

                    // Store connection
                    {
                        let mut connections = self.connections.write().await;
                        let mut domain_map = self.domain_map.write().await;

                        connections.insert(connection.id, Arc::clone(&connection));
                        domain_map.insert(connection.domain().to_string(), connection.id);
                    }

                    // Emit event
                    let _ = self.event_tx.send(ServerEvent::ModuleConnected {
                        connection_id: connection.id,
                        domain: connection.domain().to_string(),
                    });

                    log::info!(
                        "Module '{}' registered (connection {})",
                        connection.domain(),
                        connection.id
                    );

                    return Ok(connection);
                }
                Err(e) => {
                    log::warn!("Registration failed for {}: {}", addr, e);
                    continue;
                }
            }
        }
    }

    /// Perform the module registration handshake.
    async fn perform_registration(&self, stream: TcpStream) -> Result<ModuleConnection, IpcError> {
        let transport = SplitTransport::new(stream);

        // Wait for registration message
        let registration_msg = tokio::time::timeout(
            self.config.registration_timeout,
            transport.recv(),
        )
        .await
        .map_err(|_| IpcError::Timeout("Registration timeout".to_string()))??;

        // Validate registration message - check it's a Control message with "register" topic
        if registration_msg.message_type != MessageType::Control
            || !registration_msg.topic.ends_with("register")
        {
            return Err(IpcError::InvalidMessage(
                "Expected registration message".to_string(),
            ));
        }

        // Parse registration from data field
        let registration: ModuleRegistration =
            serde_json::from_value(registration_msg.data.clone())
                .map_err(|e| IpcError::InvalidMessage(format!("Invalid registration: {}", e)))?;

        // Check for duplicate domain
        {
            let domain_map = self.domain_map.read().await;
            if domain_map.contains_key(&registration.name) {
                return Err(IpcError::Connection(format!(
                    "Module '{}' already registered",
                    registration.name
                )));
            }
        }

        // Send acknowledgement - use response with same transaction_id
        let ack = registration_msg.into_response(serde_json::json!({"status": "registered"}))
            .with_topic("control.registerack");

        transport.send(&ack).await?;

        // Create connection
        let conn_id = next_connection_id();
        let (connection, outbound_rx) = ModuleConnection::new(conn_id, registration, transport);

        // Start outbound sender task
        let transport_clone = connection.clone_transport();
        let active = Arc::clone(&connection.active);
        tokio::spawn(async move {
            Self::run_outbound_sender(transport_clone, outbound_rx, active).await;
        });

        Ok(connection)
    }

    /// Task that sends outbound messages to a module.
    async fn run_outbound_sender(
        transport: SplitTransport,
        mut rx: mpsc::Receiver<CommandMessage>,
        active: Arc<AtomicBool>,
    ) {
        while active.load(Ordering::SeqCst) {
            match rx.recv().await {
                Some(msg) => {
                    if let Err(e) = transport.send(&msg).await {
                        log::error!("Failed to send to module: {}", e);
                        active.store(false, Ordering::SeqCst);
                        break;
                    }
                }
                None => break,
            }
        }
    }

    /// Get a connection by ID.
    pub async fn get_connection(&self, id: usize) -> Option<Arc<ModuleConnection>> {
        let connections = self.connections.read().await;
        connections.get(&id).cloned()
    }

    /// Get a connection by domain name.
    pub async fn get_connection_by_domain(&self, domain: &str) -> Option<Arc<ModuleConnection>> {
        let domain_map = self.domain_map.read().await;
        if let Some(&id) = domain_map.get(domain) {
            drop(domain_map);
            self.get_connection(id).await
        } else {
            None
        }
    }

    /// Send a message to a specific module by domain.
    pub async fn send_to_domain(&self, domain: &str, msg: CommandMessage) -> Result<(), IpcError> {
        let connection = self
            .get_connection_by_domain(domain)
            .await
            .ok_or_else(|| IpcError::ModuleNotFound(domain.to_string()))?;

        connection.send(msg).await
    }

    /// Broadcast a message to all connected modules.
    pub async fn broadcast(&self, msg: CommandMessage) -> Vec<Result<(), IpcError>> {
        let connections = self.connections.read().await;
        let mut results = Vec::new();

        for connection in connections.values() {
            results.push(connection.send(msg.clone()).await);
        }

        results
    }

    /// Remove a connection.
    pub async fn remove_connection(&self, id: usize) -> Option<Arc<ModuleConnection>> {
        let mut connections = self.connections.write().await;
        let mut domain_map = self.domain_map.write().await;

        if let Some(connection) = connections.remove(&id) {
            domain_map.remove(connection.domain());
            connection.deactivate();

            // Emit event
            let _ = self.event_tx.send(ServerEvent::ModuleDisconnected {
                connection_id: id,
                domain: connection.domain().to_string(),
            });

            log::info!(
                "Module '{}' disconnected (connection {})",
                connection.domain(),
                id
            );

            Some(connection)
        } else {
            None
        }
    }

    /// Get all connected module domains.
    pub async fn connected_domains(&self) -> Vec<String> {
        let domain_map = self.domain_map.read().await;
        domain_map.keys().cloned().collect()
    }

    /// Get the number of connected modules.
    pub async fn connection_count(&self) -> usize {
        let connections = self.connections.read().await;
        connections.len()
    }

    /// Subscribe to server events.
    pub fn subscribe_events(&self) -> broadcast::Receiver<ServerEvent> {
        self.event_tx.subscribe()
    }

    /// Check if the server is running.
    pub fn is_running(&self) -> bool {
        self.running.load(Ordering::SeqCst)
    }

    /// Shutdown the server.
    pub async fn shutdown(&self) {
        self.running.store(false, Ordering::SeqCst);

        // Close all connections
        let mut connections = self.connections.write().await;
        for (_, connection) in connections.drain() {
            connection.deactivate();
        }

        let mut domain_map = self.domain_map.write().await;
        domain_map.clear();
    }
}

/// Handle a module connection in a dedicated task.
///
/// This is a helper function that processes messages from a connected module
/// and dispatches them appropriately.
pub async fn handle_module_connection<F, Fut>(
    connection: Arc<ModuleConnection>,
    heartbeat_timeout: Duration,
    mut on_message: F,
) where
    F: FnMut(CommandMessage) -> Fut + Send,
    Fut: std::future::Future<Output = Option<CommandMessage>> + Send,
{
    let transport = connection.clone_transport();

    while connection.is_active() {
        // Check for heartbeat timeout
        if connection.is_timed_out(heartbeat_timeout).await {
            log::warn!(
                "Module '{}' heartbeat timeout",
                connection.domain()
            );
            connection.deactivate();
            break;
        }

        // Receive with timeout
        match tokio::time::timeout(Duration::from_secs(1), transport.recv()).await {
            Ok(Ok(msg)) => {
                // Update heartbeat on any message
                connection.update_heartbeat().await;

                // Handle heartbeat specially
                if msg.is_heartbeat() {
                    continue;
                }

                // Process message
                if let Some(response) = on_message(msg).await {
                    if let Err(e) = connection.send(response).await {
                        log::error!("Failed to send response: {}", e);
                        connection.deactivate();
                        break;
                    }
                }
            }
            Ok(Err(e)) => {
                log::error!(
                    "Error receiving from module '{}': {}",
                    connection.domain(),
                    e
                );
                connection.deactivate();
                break;
            }
            Err(_) => {
                // Timeout - continue loop to check heartbeat
                continue;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_server_config() {
        let config = IpcServerConfig::new("127.0.0.1:9200");
        assert_eq!(config.bind_addr, "127.0.0.1:9200");
        assert_eq!(config.max_connections, 64);
    }

    #[tokio::test]
    async fn test_connection_id_generation() {
        let id1 = next_connection_id();
        let id2 = next_connection_id();
        assert!(id2 > id1);
    }
}