qssh 0.0.2-alpha

//! Connection pooling for QSSH
//!
//! Manages a pool of reusable connections to reduce handshake overhead
//! and improve performance for applications making multiple connections
//! to the same servers.

use crate::{Result, QsshError, QsshConfig, client::QsshClient};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};
use tokio::time::{Duration, Instant};

/// Connection pool configuration
#[derive(Clone, Debug)]
pub struct PoolConfig {
    /// Maximum connections per host
    pub max_per_host: usize,
    /// Maximum total connections
    pub max_total: usize,
    /// Connection idle timeout
    pub idle_timeout: Duration,
    /// Maximum connection lifetime
    pub max_lifetime: Duration,
    /// Enable connection health checks
    pub health_check: bool,
    /// Health check interval
    pub health_check_interval: Duration,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            max_per_host: 5,
            max_total: 20,
            idle_timeout: Duration::from_secs(300), // 5 minutes
            max_lifetime: Duration::from_secs(3600), // 1 hour
            health_check: true,
            health_check_interval: Duration::from_secs(30),
        }
    }
}

/// Pooled connection wrapper
struct PooledConnection {
    client: QsshClient,
    created_at: Instant,
    last_used: Instant,
    use_count: usize,
    host_key: String,
}

impl PooledConnection {
    fn new(client: QsshClient, host_key: String) -> Self {
        let now = Instant::now();
        Self {
            client,
            created_at: now,
            last_used: now,
            use_count: 0,
            host_key,
        }
    }

    fn is_expired(&self, config: &PoolConfig) -> bool {
        let now = Instant::now();

        // Check lifetime
        if now.duration_since(self.created_at) > config.max_lifetime {
            return true;
        }

        // Check idle timeout
        if now.duration_since(self.last_used) > config.idle_timeout {
            return true;
        }

        false
    }

    async fn is_healthy(&self) -> bool {
        // Send a ping to check if connection is still alive
        if let Some(transport) = self.client.transport() {
            use crate::transport::Message;
            let ping = Message::Ping(rand::random());
            transport.send_message(&ping).await.is_ok()
        } else {
            false
        }
    }
}

/// Connection pool manager
pub struct ConnectionPool {
    config: PoolConfig,
    connections: Arc<RwLock<HashMap<String, Vec<Arc<Mutex<PooledConnection>>>>>>,
    total_count: Arc<Mutex<usize>>,
}

impl ConnectionPool {
    /// Create new connection pool
    pub fn new(config: PoolConfig) -> Self {
        let pool = Self {
            config: config.clone(),
            connections: Arc::new(RwLock::new(HashMap::new())),
            total_count: Arc::new(Mutex::new(0)),
        };

        // Start background tasks
        if config.health_check {
            let pool_clone = pool.clone();
            tokio::spawn(async move {
                pool_clone.health_check_task().await;
            });
        }

        // Start cleanup task
        let pool_clone = pool.clone();
        tokio::spawn(async move {
            pool_clone.cleanup_task().await;
        });

        pool
    }

    /// Get a connection from the pool
    pub async fn get(&self, config: &QsshConfig) -> Result<PoolHandle> {
        let host_key = Self::make_host_key(config);

        // Try to get existing connection
        if let Some(conn) = self.get_existing(&host_key).await? {
            return Ok(conn);
        }

        // Create new connection
        self.create_new(config, host_key).await
    }

    /// Get existing connection from pool
    async fn get_existing(&self, host_key: &str) -> Result<Option<PoolHandle>> {
        let connections = self.connections.read().await;

        if let Some(host_conns) = connections.get(host_key) {
            for conn_arc in host_conns {
                let mut conn = conn_arc.lock().await;

                // Skip if expired
                if conn.is_expired(&self.config) {
                    continue;
                }

                // Check health if needed
                if self.config.health_check && !conn.is_healthy().await {
                    continue;
                }

                // Found a good connection
                conn.last_used = Instant::now();
                conn.use_count += 1;

                return Ok(Some(PoolHandle {
                    pool: self.clone(),
                    connection: conn_arc.clone(),
                }));
            }
        }

        Ok(None)
    }

    /// Create new connection
    async fn create_new(&self, config: &QsshConfig, host_key: String) -> Result<PoolHandle> {
        // Check total limit
        let total = *self.total_count.lock().await;
        if total >= self.config.max_total {
            return Err(QsshError::Connection("Connection pool limit reached".into()));
        }

        // Check per-host limit
        {
            let connections = self.connections.read().await;
            if let Some(host_conns) = connections.get(&host_key) {
                if host_conns.len() >= self.config.max_per_host {
                    return Err(QsshError::Connection(
                        format!("Per-host connection limit reached for {}", host_key)
                    ));
                }
            }
        }

        // Create new client and connect
        let mut client = QsshClient::new(config.clone());
        client.connect().await?;

        // Add to pool
        let pooled = Arc::new(Mutex::new(PooledConnection::new(client, host_key.clone())));

        {
            let mut connections = self.connections.write().await;
            connections.entry(host_key.clone())
                .or_insert_with(Vec::new)
                .push(pooled.clone());
        }

        // Update count
        {
            let mut total = self.total_count.lock().await;
            *total += 1;
        }

        Ok(PoolHandle {
            pool: self.clone(),
            connection: pooled,
        })
    }

    /// Make host key from config
    fn make_host_key(config: &QsshConfig) -> String {
        format!("{}@{}", config.username, config.server)
    }

    /// Background health check task
    async fn health_check_task(&self) {
        let mut interval = tokio::time::interval(self.config.health_check_interval);

        loop {
            interval.tick().await;

            let connections = self.connections.read().await;
            for (_host, conns) in connections.iter() {
                for conn_arc in conns {
                    let conn = conn_arc.lock().await;
                    if !conn.is_healthy().await {
                        log::warn!("Unhealthy connection detected for {}", conn.host_key);
                        // Mark for removal in cleanup
                    }
                }
            }
        }
    }

    /// Background cleanup task
    async fn cleanup_task(&self) {
        let mut interval = tokio::time::interval(Duration::from_secs(60));

        loop {
            interval.tick().await;

            let mut connections = self.connections.write().await;
            let mut total_removed = 0;

            // Clean up expired connections
            for (_host, conns) in connections.iter_mut() {
                conns.retain(|conn_arc| {
                    if let Ok(conn) = conn_arc.try_lock() {
                        if conn.is_expired(&self.config) {
                            log::debug!("Removing expired connection for {}", conn.host_key);
                            total_removed += 1;
                            return false;
                        }
                    }
                    true
                });
            }

            // Remove empty entries
            connections.retain(|_, conns| !conns.is_empty());

            // Update total count
            if total_removed > 0 {
                let mut total = self.total_count.lock().await;
                *total = (*total).saturating_sub(total_removed);
            }
        }
    }

    /// Get pool statistics
    pub async fn stats(&self) -> PoolStats {
        let connections = self.connections.read().await;
        let total = *self.total_count.lock().await;

        let mut stats = PoolStats {
            total_connections: total,
            connections_per_host: HashMap::new(),
            total_use_count: 0,
            expired_count: 0,
        };

        for (host, conns) in connections.iter() {
            stats.connections_per_host.insert(host.clone(), conns.len());

            for conn_arc in conns {
                if let Ok(conn) = conn_arc.try_lock() {
                    stats.total_use_count += conn.use_count;
                    if conn.is_expired(&self.config) {
                        stats.expired_count += 1;
                    }
                }
            }
        }

        stats
    }

    /// Clear all connections
    pub async fn clear(&self) {
        let mut connections = self.connections.write().await;
        connections.clear();

        let mut total = self.total_count.lock().await;
        *total = 0;
    }
}

impl Clone for ConnectionPool {
    fn clone(&self) -> Self {
        Self {
            config: self.config.clone(),
            connections: self.connections.clone(),
            total_count: self.total_count.clone(),
        }
    }
}

/// Handle to a pooled connection
pub struct PoolHandle {
    pool: ConnectionPool,
    connection: Arc<Mutex<PooledConnection>>,
}

impl PoolHandle {
    /// Execute a function with the client
    pub async fn with_client<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut QsshClient) -> R,
    {
        let mut conn = self.connection.lock().await;
        f(&mut conn.client)
    }

    /// Return connection to pool (happens automatically on drop)
    pub fn release(self) {
        // Connection is automatically returned when PoolHandle is dropped
        drop(self);
    }
}

/// Pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
    pub total_connections: usize,
    pub connections_per_host: HashMap<String, usize>,
    pub total_use_count: usize,
    pub expired_count: usize,
}

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

    #[tokio::test]
    async fn test_pool_config_default() {
        let config = PoolConfig::default();
        assert_eq!(config.max_per_host, 5);
        assert_eq!(config.max_total, 20);
        assert_eq!(config.idle_timeout, Duration::from_secs(300));
    }

    #[tokio::test]
    async fn test_host_key_generation() {
        let config = QsshConfig {
            server: "test.example.com:22".to_string(),
            username: "alice".to_string(),
            password: None,
            port_forwards: vec![],
            use_qkd: false,
            pq_algorithm: crate::PqAlgorithm::Falcon512,
            key_rotation_interval: 3600,
        };

        let host_key = ConnectionPool::make_host_key(&config);
        assert_eq!(host_key, "alice@test.example.com:22");
    }

    #[tokio::test]
    async fn test_pool_stats() {
        let pool = ConnectionPool::new(PoolConfig::default());
        let stats = pool.stats().await;

        assert_eq!(stats.total_connections, 0);
        assert_eq!(stats.total_use_count, 0);
        assert_eq!(stats.expired_count, 0);
    }
}