cnctd-service-ssh 0.1.8

//! Session registry for managing interactive shell sessions.

use crate::service_error::ServiceError;
use crate::sessions::connection::SshConnection;
use crate::sessions::terminal::TerminalEmulator;
use crate::sessions::types::*;
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, warn};

/// Maximum output buffer size (1MB)
const MAX_BUFFER_SIZE: usize = 1024 * 1024;

/// Output buffer with ring buffer behavior
pub struct OutputBuffer {
    data: Vec<u8>,
    truncated: bool,
}

impl OutputBuffer {
    fn new() -> Self {
        Self {
            data: Vec::new(),
            truncated: false,
        }
    }

    fn append(&mut self, new_data: &[u8]) {
        self.data.extend_from_slice(new_data);

        // Truncate oldest data if buffer exceeds max size
        if self.data.len() > MAX_BUFFER_SIZE {
            let excess = self.data.len() - MAX_BUFFER_SIZE;
            self.data.drain(0..excess);
            self.truncated = true;
        }
    }

    fn take(&mut self) -> (Vec<u8>, bool) {
        let truncated = self.truncated;
        self.truncated = false;
        (std::mem::take(&mut self.data), truncated)
    }

    fn peek(&self) -> (&[u8], bool) {
        (&self.data, self.truncated)
    }

    fn len(&self) -> usize {
        self.data.len()
    }
}

/// Active shell session
pub struct ShellSession {
    /// Unique session ID
    pub id: String,
    /// SSH target ID
    pub target_id: String,
    /// Human-readable name
    pub name: Option<String>,
    /// Client identifier
    pub client_id: Option<String>,
    /// Remote tmux session name
    pub tmux_session: String,
    /// Terminal dimensions
    pub cols: u16,
    pub rows: u16,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
    /// Last activity timestamp
    pub last_activity: Arc<RwLock<DateTime<Utc>>>,
    /// Session state
    pub state: Arc<RwLock<SessionState>>,
    /// Output buffer
    buffer: Arc<RwLock<OutputBuffer>>,
    /// Terminal emulator
    terminal: Arc<RwLock<TerminalEmulator>>,
    /// SSH connection (None if disconnected)
    connection: Arc<RwLock<Option<SshConnection>>>,
    /// Shutdown signal for output reader task
    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
}

impl ShellSession {
    /// Create a new session
    pub fn new(
        id: String,
        target_id: String,
        name: Option<String>,
        client_id: Option<String>,
        tmux_session: String,
        cols: u16,
        rows: u16,
        connection: SshConnection,
    ) -> Self {
        let now = Utc::now();
        Self {
            id,
            target_id,
            name,
            client_id,
            tmux_session,
            cols,
            rows,
            created_at: now,
            last_activity: Arc::new(RwLock::new(now)),
            state: Arc::new(RwLock::new(SessionState::Active)),
            buffer: Arc::new(RwLock::new(OutputBuffer::new())),
            terminal: Arc::new(RwLock::new(TerminalEmulator::new(cols, rows))),
            connection: Arc::new(RwLock::new(Some(connection))),
            shutdown_tx: None,
        }
    }

    /// Start the background output reader task
    pub fn start_output_reader(&mut self) {
        let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
        self.shutdown_tx = Some(shutdown_tx);

        let connection = Arc::clone(&self.connection);
        let buffer = Arc::clone(&self.buffer);
        let terminal = Arc::clone(&self.terminal);
        let state = Arc::clone(&self.state);
        let last_activity = Arc::clone(&self.last_activity);
        let session_id = self.id.clone();

        tokio::spawn(async move {
            debug!("Starting output reader for session {}", session_id);

            loop {
                // Check for shutdown signal
                if shutdown_rx.try_recv().is_ok() {
                    debug!("Output reader shutting down for session {}", session_id);
                    break;
                }

                // Try to read from connection
                let read_result = {
                    let mut conn_guard = connection.write().await;
                    if let Some(conn) = conn_guard.as_mut() {
                        // Check if connection is alive
                        if !conn.is_alive() {
                            warn!("Connection died for session {}", session_id);
                            *state.write().await = SessionState::Disconnected;
                            *conn_guard = None;
                            None
                        } else {
                            conn.try_recv()
                        }
                    } else {
                        None
                    }
                };

                if let Some(data) = read_result {
                    if !data.is_empty() {
                        // Update last activity
                        *last_activity.write().await = Utc::now();

                        // Append to buffer
                        buffer.write().await.append(&data);

                        // Process through terminal emulator
                        terminal.read().await.process(&data).await;
                    }
                } else {
                    // Check if connection is still present
                    let conn_guard = connection.read().await;
                    if conn_guard.is_none() {
                        debug!("Connection removed for session {}", session_id);
                        break;
                    }
                }

                // Small sleep to prevent busy loop
                tokio::time::sleep(std::time::Duration::from_millis(20)).await;
            }

            debug!("Output reader finished for session {}", session_id);
        });
    }

    /// Get session info
    pub async fn info(&self) -> ShellSessionInfo {
        ShellSessionInfo {
            id: self.id.clone(),
            target_id: self.target_id.clone(),
            name: self.name.clone(),
            client_id: self.client_id.clone(),
            state: *self.state.read().await,
            created_at: self.created_at.to_rfc3339(),
            last_activity: self.last_activity.read().await.to_rfc3339(),
            size: (self.cols, self.rows),
        }
    }

    /// Write data to the session
    pub async fn write(&self, data: &[u8]) -> Result<usize, ServiceError> {
        let mut conn_guard = self.connection.write().await;
        if let Some(conn) = conn_guard.as_mut() {
            conn.send(data).await?;
            *self.last_activity.write().await = Utc::now();
            Ok(data.len())
        } else {
            Err(ServiceError::Internal(
                "Session is disconnected".to_string(),
            ))
        }
    }

    /// Read from buffer
    pub async fn read(&self, consume: bool) -> (String, usize, bool) {
        let mut buffer = self.buffer.write().await;
        if consume {
            let (data, truncated) = buffer.take();
            let text = String::from_utf8_lossy(&data).to_string();
            (text, 0, truncated)
        } else {
            let (data, truncated) = buffer.peek();
            let text = String::from_utf8_lossy(data).to_string();
            (text, data.len(), truncated)
        }
    }

    /// Get current screen state
    pub async fn screen_state(&self) -> ScreenState {
        self.terminal.read().await.screen_state().await
    }

    /// Resize the terminal
    pub async fn resize(&mut self, cols: u16, rows: u16) -> Result<(), ServiceError> {
        // Resize local terminal emulator
        self.terminal.write().await.resize(cols, rows).await;

        // Resize remote PTY
        let mut conn_guard = self.connection.write().await;
        if let Some(conn) = conn_guard.as_mut() {
            conn.resize(cols, rows)?;
        }

        self.cols = cols;
        self.rows = rows;

        Ok(())
    }

    /// Wait for output with timeout
    pub async fn wait_for_output(&self, timeout_ms: u64, min_bytes: usize) -> bool {
        let start = std::time::Instant::now();
        let timeout = std::time::Duration::from_millis(timeout_ms);

        while start.elapsed() < timeout {
            if self.buffer.read().await.len() >= min_bytes {
                return true;
            }
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        }

        self.buffer.read().await.len() >= min_bytes
    }

    /// Get the screen content as a single string (for pattern matching)
    pub async fn get_screen_text(&self) -> String {
        let screen_state = self.terminal.read().await.screen_state().await;
        screen_state.lines.join("\n")
    }

    /// Wait for a pattern to appear in the screen output
    /// Returns true if pattern found, false if timeout
    pub async fn wait_for_pattern(&self, pattern: &str, timeout_ms: u64) -> bool {
        let start = std::time::Instant::now();
        let timeout = std::time::Duration::from_millis(timeout_ms);

        while start.elapsed() < timeout {
            let screen_text = self.get_screen_text().await;
            if screen_text.contains(pattern) {
                return true;
            }
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        }

        // Final check
        self.get_screen_text().await.contains(pattern)
    }

    /// Wait for the screen output to stabilize (no changes for stable_ms)
    /// Returns true if stabilized, false if timeout
    pub async fn wait_for_stable(&self, stable_ms: u64, timeout_ms: u64) -> bool {
        let start = std::time::Instant::now();
        let timeout = std::time::Duration::from_millis(timeout_ms);
        let stable_duration = std::time::Duration::from_millis(stable_ms);

        let mut last_screen = self.get_screen_text().await;
        let mut last_change = std::time::Instant::now();

        while start.elapsed() < timeout {
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;

            let current_screen = self.get_screen_text().await;
            if current_screen != last_screen {
                last_screen = current_screen;
                last_change = std::time::Instant::now();
            } else if last_change.elapsed() >= stable_duration {
                return true;
            }
        }

        false
    }

    /// Close the session
    pub async fn close(mut self, _force: bool) -> Result<bool, ServiceError> {
        // Signal output reader to stop
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            let _ = shutdown_tx.send(());
        }

        // Close the SSH connection
        let connection = self.connection.write().await.take();
        if let Some(conn) = connection {
            conn.close().await?;
        }

        *self.state.write().await = SessionState::Closed;

        Ok(true)
    }
}

/// Registry for managing shell sessions
pub struct ShellSessionRegistry {
    sessions: Arc<RwLock<HashMap<String, Arc<RwLock<ShellSession>>>>>,
}

impl ShellSessionRegistry {
    /// Create a new registry
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Add a session to the registry
    pub async fn add(&self, session: ShellSession) {
        let id = session.id.clone();
        let session = Arc::new(RwLock::new(session));
        self.sessions.write().await.insert(id, session);
    }

    /// Get a session by ID
    pub async fn get(&self, id: &str) -> Option<Arc<RwLock<ShellSession>>> {
        self.sessions.read().await.get(id).cloned()
    }

    /// Remove a session by ID
    pub async fn remove(&self, id: &str) -> Option<Arc<RwLock<ShellSession>>> {
        self.sessions.write().await.remove(id)
    }

    /// List sessions with optional filters
    pub async fn list(
        &self,
        target_id: Option<&str>,
        client_id: Option<&str>,
        include_disconnected: bool,
    ) -> Vec<ShellSessionInfo> {
        let sessions = self.sessions.read().await;
        let mut result = Vec::new();

        for session_lock in sessions.values() {
            let session = session_lock.read().await;
            let state = *session.state.read().await;

            // Filter by state
            if !include_disconnected && state == SessionState::Disconnected {
                continue;
            }

            // Filter by target_id
            if let Some(tid) = target_id {
                if session.target_id != tid {
                    continue;
                }
            }

            // Filter by client_id
            if let Some(cid) = client_id {
                match &session.client_id {
                    Some(session_cid) if session_cid == cid => {}
                    _ => continue,
                }
            }

            result.push(session.info().await);
        }

        result
    }

    /// Get count of active sessions
    pub async fn active_count(&self) -> usize {
        let sessions = self.sessions.read().await;
        let mut count = 0;
        for session_lock in sessions.values() {
            let session = session_lock.read().await;
            if *session.state.read().await == SessionState::Active {
                count += 1;
            }
        }
        count
    }
}

impl Default for ShellSessionRegistry {
    fn default() -> Self {
        Self::new()
    }
}