use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::Result;
use russh::client;
use sha2::{Digest, Sha256};
use tokio::sync::{Mutex, RwLock};
use tracing;

use super::super::backend::russh::connect::{connect_and_auth, ClientHandler};
use super::super::backend::russh::exec::{disconnect_session, exec_collect};
use super::super::models::config::{AuthMethod, SshConfig};
use super::super::models::types::{ExecRequest, ExecResult};
use super::super::shell::detect_remote_shell;

/// Connection pool configuration.
#[derive(Clone, Debug)]
pub struct PoolConfig {
    /// Maximum number of idle connections.
    pub max_idle: usize,
    /// Maximum idle time for a connection.
    pub max_idle_time: Duration,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            max_idle: 10,
            max_idle_time: Duration::from_secs(300),
        }
    }
}

/// Persistent SSH session.
///
/// Holds an authenticated SSH connection that can be reused for multiple operations without reconnecting each time.
/// Suitable for automation scripts that need to execute multiple commands sequentially.
pub struct Session {
    cfg: SshConfig,
    pub(crate) handle: client::Handle<ClientHandler>,
}

impl Session {
    /// Establish a new persistent session.
    pub async fn connect(cfg: SshConfig) -> Result<Self> {
        let handle = connect_and_auth(cfg.clone()).await?;
        Ok(Self { cfg, handle })
    }

    /// Construct a Session from an existing handle (for connection pool internal management).
    pub(crate) fn from_handle(cfg: SshConfig, handle: client::Handle<ClientHandler>) -> Self {
        Self { cfg, handle }
    }

    /// Get the session configuration.
    pub fn config(&self) -> &SshConfig {
        &self.cfg
    }

    /// Execute a command.
    pub async fn exec(&self, req: ExecRequest) -> Result<ExecResult> {
        let shell = detect_remote_shell(&self.handle).await;
        let command = shell.build_exec_command(&req);

        let result = if let Some(timeout_ms) = req.timeout_ms {
            match tokio::time::timeout(
                Duration::from_millis(timeout_ms),
                exec_collect(&self.handle, &command),
            )
            .await
            {
                Ok(exec_result) => {
                    let (exit_code, stdout, stderr) = exec_result?;
                    ExecResult {
                        exit_code,
                        stdout,
                        stderr,
                        timed_out: false,
                    }
                }
                Err(_) => {
                    tracing::warn!(command = %command, "exec command timed out");
                    ExecResult {
                        exit_code: 124,
                        stdout: Vec::new(),
                        stderr: b"command timed out".to_vec(),
                        timed_out: true,
                    }
                }
            }
        } else {
            let (exit_code, stdout, stderr) = exec_collect(&self.handle, &command).await?;
            ExecResult {
                exit_code,
                stdout,
                stderr,
                timed_out: false,
            }
        };

        Ok(result)
    }

    /// Close the session.
    pub async fn close(self) -> Result<()> {
        disconnect_session(&self.handle).await
    }
}

/// Connection pool.
///
/// Caches authenticated SSH connections to avoid re-establishing connections for each operation.
#[derive(Clone)]
pub struct ConnectionPool {
    pool: Arc<RwLock<HashMap<String, Arc<Mutex<Option<PoolEntry>>>>>>,
    config: PoolConfig,
}

impl std::fmt::Debug for ConnectionPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConnectionPool")
            .field("config", &self.config)
            .finish()
    }
}

struct PoolEntry {
    handle: client::Handle<ClientHandler>,
    last_used: Instant,
}

impl ConnectionPool {
    /// Create a new connection pool.
    pub fn new(config: PoolConfig) -> Self {
        Self {
            pool: Arc::new(RwLock::new(HashMap::new())),
            config,
        }
    }

    /// Create a connection pool with default configuration.
    pub fn default_pool() -> Self {
        Self::new(PoolConfig::default())
    }

    /// Get the pool key.
    ///
    /// Uses SHA256 hash for password auth, includes path and passphrase hash for private key auth,
    /// to prevent different credentials from sharing a connection.
    fn pool_key(cfg: &SshConfig) -> String {
        let auth_suffix = match &cfg.auth {
            AuthMethod::Password { password } => {
                let mut hasher = Sha256::new();
                hasher.update(password.as_bytes());
                format!("pw:{:x}", hasher.finalize())
            }
            AuthMethod::PrivateKey {
                private_key_path,
                passphrase,
            } => {
                let pass_hash = passphrase
                    .as_ref()
                    .map(|p| {
                        let mut h = Sha256::new();
                        h.update(p.as_bytes());
                        format!(":{:x}", h.finalize())
                    })
                    .unwrap_or_default();
                format!("key:{}{}", private_key_path, pass_hash)
            }
        };
        format!("{}@{}:{}:{}", cfg.user, cfg.host, cfg.port, auth_suffix)
    }

    /// Get or create a connection from the pool.
    ///
    /// Uses per-key mutex to prevent concurrent connection creation for the same key.
    pub async fn get_or_connect(&self, cfg: &SshConfig) -> Result<Session> {
        let key = Self::pool_key(cfg);

        // Get or create the per-key lock for this key
        let entry_lock = {
            let mut pool = self.pool.write().await;
            pool.entry(key.clone())
                .or_insert_with(|| Arc::new(Mutex::new(None)))
                .clone()
        };

        // Check/create connection within the per-key lock
        let mut guard = entry_lock.lock().await;
        if let Some(entry) = guard.take() {
            // Check if expired
            if entry.last_used.elapsed() < self.config.max_idle_time {
                tracing::debug!(key = %key, "reusing pooled connection");
                return Ok(Session {
                    cfg: cfg.clone(),
                    handle: entry.handle,
                });
            }
            // Expired, explicitly disconnect
            tracing::debug!(key = %key, "pooled connection expired, creating new");
            tokio::spawn(async move {
                let _ = disconnect_session(&entry.handle).await;
            });
        }
        // guard is None after take, clean up empty slot
        drop(guard);
        {
            let mut pool = self.pool.write().await;
            // Only remove if entry_lock still points to the same Arc (avoid removing a newly created one)
            if let Some(lock) = pool.get(&key) {
                if Arc::ptr_eq(lock, &entry_lock) {
                    pool.remove(&key);
                }
            }
        }

        // Create a new connection
        tracing::debug!(key = %key, "creating new connection");
        Session::connect(cfg.clone()).await
    }

    /// Return a connection to the pool.
    pub async fn return_connection(&self, session: Session) {
        let key = Self::pool_key(&session.cfg);

        let mut pool = self.pool.write().await;

        // Clean up empty slots and check actual connection count
        pool.retain(|_, lock| {
            // Synchronous check: if try_lock succeeds and is None, remove the slot
            lock.try_lock().map_or(true, |guard| guard.is_some())
        });

        // Check pool size limit, remove the oldest connection
        if pool.len() >= self.config.max_idle {
            let mut oldest_key: Option<String> = None;
            let mut oldest_time: Option<Instant> = None;
            for (k, lock) in pool.iter() {
                if let Ok(guard) = lock.try_lock() {
                    if let Some(ref entry) = *guard {
                        if oldest_time.is_none() || entry.last_used < oldest_time.unwrap() {
                            oldest_time = Some(entry.last_used);
                            oldest_key = Some(k.clone());
                        }
                    }
                }
            }
            if let Some(old_key) = oldest_key {
                if let Some(entry_lock) = pool.remove(&old_key) {
                    // Use lock().await instead of try_lock() to ensure the connection is definitely disconnected
                    let mut guard = entry_lock.lock().await;
                    if let Some(entry) = guard.take() {
                        tokio::spawn(async move {
                            let _ = disconnect_session(&entry.handle).await;
                        });
                    }
                }
            }
        }

        let entry_lock = pool
            .entry(key)
            .or_insert_with(|| Arc::new(Mutex::new(None)));
        let mut guard = entry_lock.lock().await;
        *guard = Some(PoolEntry {
            handle: session.handle,
            last_used: Instant::now(),
        });
    }

    /// Clean up expired connections.
    ///
    /// Collects and removes expired connections in a single write lock to avoid TOCTOU race conditions.
    pub async fn cleanup(&self) {
        let mut pool = self.pool.write().await;
        let now = Instant::now();
        let expired: Vec<_> = pool
            .iter()
            .filter_map(|(key, lock)| {
                if let Ok(guard) = lock.try_lock() {
                    if let Some(ref entry) = *guard {
                        if now.duration_since(entry.last_used) >= self.config.max_idle_time {
                            return Some(key.clone());
                        }
                    }
                }
                None
            })
            .collect();

        for key in expired {
            if let Some(entry_lock) = pool.remove(&key) {
                let mut guard = entry_lock.lock().await;
                if let Some(entry) = guard.take() {
                    tokio::spawn(async move {
                        let _ = disconnect_session(&entry.handle).await;
                    });
                }
            }
        }
    }

    /// Clear the connection pool.
    pub async fn clear(&self) {
        let mut pool = self.pool.write().await;
        for (_, entry_lock) in pool.drain() {
            let mut guard = entry_lock.lock().await;
            if let Some(entry) = guard.take() {
                tokio::spawn(async move {
                    let _ = disconnect_session(&entry.handle).await;
                });
            }
        }
    }

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

    /// Check if the pool is empty.
    pub async fn is_empty(&self) -> bool {
        let pool = self.pool.read().await;
        pool.is_empty()
    }
}

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

    fn password_cfg(password: &str) -> SshConfig {
        SshConfig {
            host: "10.0.0.1".to_string(),
            port: 22,
            user: "root".to_string(),
            auth: AuthMethod::Password {
                password: password.to_string(),
            },
            timeout_ms: 5000,
            host_key_policy: Default::default(),
            known_hosts_path: None,
        }
    }

    fn key_cfg(path: &str, passphrase: Option<&str>) -> SshConfig {
        SshConfig {
            host: "10.0.0.1".to_string(),
            port: 22,
            user: "root".to_string(),
            auth: AuthMethod::PrivateKey {
                private_key_path: path.to_string(),
                passphrase: passphrase.map(|s| s.to_string()),
            },
            timeout_ms: 5000,
            host_key_policy: Default::default(),
            known_hosts_path: None,
        }
    }

    // ===== pool_key security tests =====

    #[test]
    fn pool_key_different_passwords_produce_different_keys() {
        let key_a = ConnectionPool::pool_key(&password_cfg("alpha"));
        let key_b = ConnectionPool::pool_key(&password_cfg("beta"));
        assert_ne!(
            key_a, key_b,
            "different passwords must produce different pool keys"
        );
    }

    #[test]
    fn pool_key_same_length_passwords_produce_different_keys() {
        // Core security fix: old code used password.len(), same length would collide
        let key_a = ConnectionPool::pool_key(&password_cfg("abc"));
        let key_b = ConnectionPool::pool_key(&password_cfg("xyz"));
        assert_ne!(
            key_a, key_b,
            "same-length but different passwords must produce different pool keys"
        );
    }

    #[test]
    fn pool_key_same_password_produces_same_key() {
        let key_a = ConnectionPool::pool_key(&password_cfg("secret"));
        let key_b = ConnectionPool::pool_key(&password_cfg("secret"));
        assert_eq!(key_a, key_b, "same password must produce the same pool key");
    }

    #[test]
    fn pool_key_password_hash_does_not_contain_plaintext() {
        let key = ConnectionPool::pool_key(&password_cfg("my_secret_password"));
        assert!(
            !key.contains("my_secret_password"),
            "pool key should not contain plaintext password"
        );
    }

    #[test]
    fn pool_key_different_passphrases_produce_different_keys() {
        let key_a = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_rsa", Some("alpha")));
        let key_b = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_rsa", Some("beta")));
        assert_ne!(
            key_a, key_b,
            "same key path but different passphrases must produce different pool keys"
        );
    }

    #[test]
    fn pool_key_none_passphrase_differs_from_empty_passphrase() {
        let key_none = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_rsa", None));
        let key_empty = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_rsa", Some("")));
        assert_ne!(key_none, key_empty);
    }

    #[test]
    fn pool_key_different_paths_produce_different_keys() {
        let key_a = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_rsa", None));
        let key_b = ConnectionPool::pool_key(&key_cfg("/root/.ssh/id_ed25519", None));
        assert_ne!(key_a, key_b);
    }

    #[test]
    fn pool_key_different_hosts_produce_different_keys() {
        let mut cfg_a = password_cfg("pass");
        cfg_a.host = "host-a".to_string();
        let mut cfg_b = password_cfg("pass");
        cfg_b.host = "host-b".to_string();
        assert_ne!(
            ConnectionPool::pool_key(&cfg_a),
            ConnectionPool::pool_key(&cfg_b)
        );
    }

    #[test]
    fn pool_key_different_ports_produce_different_keys() {
        let mut cfg_a = password_cfg("pass");
        cfg_a.port = 22;
        let mut cfg_b = password_cfg("pass");
        cfg_b.port = 2222;
        assert_ne!(
            ConnectionPool::pool_key(&cfg_a),
            ConnectionPool::pool_key(&cfg_b)
        );
    }

    #[test]
    fn pool_key_different_users_produce_different_keys() {
        let mut cfg_a = password_cfg("pass");
        cfg_a.user = "root".to_string();
        let mut cfg_b = password_cfg("pass");
        cfg_b.user = "admin".to_string();
        assert_ne!(
            ConnectionPool::pool_key(&cfg_a),
            ConnectionPool::pool_key(&cfg_b)
        );
    }

    // ===== PoolConfig tests =====

    #[test]
    fn pool_config_default_values() {
        let config = PoolConfig::default();
        assert_eq!(config.max_idle, 10);
        assert_eq!(config.max_idle_time, Duration::from_secs(300));
    }

    // ===== ConnectionPool basic operations tests =====

    #[tokio::test]
    async fn pool_starts_empty() {
        let pool = ConnectionPool::default_pool();
        assert!(pool.is_empty().await);
        assert_eq!(pool.len().await, 0);
    }

    #[tokio::test]
    async fn pool_clear_on_empty_pool() {
        let pool = ConnectionPool::default_pool();
        pool.clear().await; // should not panic
        assert!(pool.is_empty().await);
    }
}
