//! SSH 连接池（P2 连接复用）。
//!
//! # 语义
//!
//! 同一「身份」（用户名 + 主机 + 端口 + 认证材料指纹）的多个会话
//! **共享同一条 SSH 连接**：第二个会话跳过 TCP / 握手 / 认证，直接在
//! 既有连接上开新的会话通道。与 OpenSSH ControlMaster 的行为一致：
//! - 打开：首个会话建立连接；
//! - 共享：后续会话引用计数 +1；
//! - 关闭：会话关闭只减引用并关闭**自己的 shell 通道**；
//!   **最后一个引用释放时**才断开底层连接（含跳板机链）。
//!
//! # 为什么 Handle 必须经由池共享
//!
//! `russh::client::Handle` 不实现 `Clone`（内含 session actor 的接收端），
//! 此前每个 `SshSession` 独占一个 Handle，无从复用。池条目持有
//! `Arc<tokio::sync::Mutex<Option<Handle>>>` 槽位，会话间共享同一 Arc；
//! tokio Mutex（而非 std）是为了允许**连接建立期间跨 `.await` 持锁**——
//! 它天然串行化了「双击两个标签同时连同一主机」的竞态：后到者等待，
//! 先到者成功后直接复用。
//!
//! # 跳板机链的归属
//!
//! 经跳板链建立的连接，其跳板 Handle 挂在**池条目**上而不是首建会话上：
//! 否则首建会话关闭时会连带剪断仍在被其他会话使用的隧道。
//! 最后一个引用释放时，跳板与目标连接一起断开。
//!
//! # 已知取舍
//!
//! - 共享连接的 keepalive / 加密参数取自**首个**建立它的会话；
//! - 共享连接断开（网络故障）时，挂在上面的所有会话一起进入 Closed——
//!   这与「它们本来就在同一条 TCP 上」的物理事实一致；
//! - 远程转发（-R）的入站路由按端口做**全局**匹配（见 forward 模块），
//!   因为入站通道总是从属连接级 Handler，而 Handler 的 session_id 属于首建会话。
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use russh::client::Handle;
use russh::Disconnect;

use super::SshHandler;

/// 可克隆的池句柄（`TerminalManager` 持有一份，每个会话 clone 一份）。
#[derive(Clone, Default)]
pub struct ConnectionPool {
    conns: Arc<Mutex<HashMap<String, PooledEntry>>>,
}

struct PooledEntry {
    /// 共享槽位：`None` = 连接建立中（或失败）；`Some` = 已就绪。
    slot: Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>>,
    /// 仍在使用此连接的会话数
    refcount: usize,
    /// 连接是否已就绪（同步可查；`has_handle` 用）
    ready: bool,
    /// 跳板机链的连接（经跳板建立时非空；随条目共享，最后释放时断开）
    hops: Vec<Handle<SshHandler>>,
    /// 日志用描述
    label: String,
}

impl ConnectionPool {
    /// 取（或创建）某身份的连接槽位，引用计数 +1。
    ///
    /// 返回的 Arc 就是池条目里的槽位本身：会话把它存进 `SshInner.handle`，
    /// 连接建立后写 `Some(handle)`，同键的其他会话即刻可见。
    pub fn slot(&self, key: &str) -> Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>> {
        let mut conns = self.conns.lock().unwrap_or_else(|e| e.into_inner());
        let entry = conns.entry(key.to_string()).or_insert_with(|| PooledEntry {
            slot: Arc::new(tokio::sync::Mutex::new(None)),
            refcount: 0,
            ready: false,
            hops: Vec::new(),
            label: String::new(),
        });
        entry.refcount += 1;
        entry.slot.clone()
    }

    /// 设置日志用描述（连接建立成功后调用）。
    pub fn set_label(&self, key: &str, label: &str) {
        if let Some(e) = self
            .conns
            .lock()
            .unwrap_or_else(|x| x.into_inner())
            .get_mut(key)
        {
            e.label = label.to_string();
        }
    }

    /// 标记连接已就绪（do_connect 写入 Handle 之后）。
    pub fn mark_ready(&self, key: &str) {
        if let Some(e) = self
            .conns
            .lock()
            .unwrap_or_else(|x| x.into_inner())
            .get_mut(key)
        {
            e.ready = true;
        }
    }

    /// 把跳板机连接挂到池条目上（fresh 连接路径、有跳板时调用一次）。
    pub fn attach_hops(&self, key: &str, hops: Vec<Handle<SshHandler>>) {
        if let Some(e) = self
            .conns
            .lock()
            .unwrap_or_else(|x| x.into_inner())
            .get_mut(key)
        {
            e.hops = hops;
        }
    }

    /// 连接是否已就绪（同步可查；替代原 `SshSession::has_handle` 的语义）。
    pub fn is_ready(&self, key: &str) -> bool {
        self.conns
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .get(key)
            .is_some_and(|e| e.ready)
    }

    /// 释放一个会话的引用。
    ///
    /// 返回 `Some((槽位, 跳板连接))` 表示这是**最后一个**引用——调用方负责
    /// 断开底层连接与跳板（异步任务里做，见 `kill`）。非最后引用返回 `None`，
    /// 调用方只需关闭自己的 shell 通道。
    pub fn release(
        &self,
        key: &str,
    ) -> Option<(
        Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>>,
        Vec<Handle<SshHandler>>,
    )> {
        let mut conns = self.conns.lock().unwrap_or_else(|e| e.into_inner());
        let Some(entry) = conns.get_mut(key) else {
            return None;
        };
        entry.refcount = entry.refcount.saturating_sub(1);
        if entry.refcount > 0 {
            return None;
        }
        // 最后一个引用：移除条目并交出断开责任
        let entry = conns.remove(key)?;
        Some((entry.slot, entry.hops))
    }
}

/// 连接池身份键：用户名 + 主机 + 端口 + 认证方式 + 认证材料指纹。
///
/// 认证材料（密码或私钥文本）取短哈希入键——同一主机配置两份不同密钥/密码时
/// 不应共享连接（那等于用 A 的身份看了 B 的会话）。
pub fn pool_key_of(
    username: &str,
    host: &str,
    port: u16,
    auth_method: &str,
    auth_material: Option<&str>,
) -> String {
    let marker = match auth_material {
        Some(m) => short_hash(m),
        None => "none".to_string(),
    };
    format!("{username}|{host}:{port}|{auth_method}|{marker}")
}

/// 材料指纹：SHA-256 前 8 字节的十六进制（16 字符）。
fn short_hash(material: &str) -> String {
    use sha2::{Digest, Sha256};
    let digest = Sha256::digest(material.as_bytes());
    digest[..8].iter().map(|b| format!("{b:02x}")).collect()
}

/// 断开一个 Handle（kill 与最后引用释放共用的收尾动作）。
pub async fn disconnect(handle: Handle<SshHandler>) {
    let _ = handle
        .disconnect(Disconnect::ByApplication, "closed by user", "")
        .await;
}

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

    /// 条目级语义测试：不涉及真实 Handle（槽位保持 None 即可）。
    #[test]
    fn acquire_increments_and_release_removes_on_last() {
        let pool = ConnectionPool::default();
        let s1 = pool.slot("k");
        let s2 = pool.slot("k");
        // 同键两次 acquire 返回同一个 Arc（这才是「共享」）
        assert!(Arc::ptr_eq(&s1, &s2));
        assert!(pool.release("k").is_none(), "还有 1 个引用，不应触发拆除");
        let (slot, hops) = pool.release("k").expect("最后一个引用应触发拆除");
        assert!(Arc::ptr_eq(&slot, &s1));
        assert!(hops.is_empty());
        // 移除后再次 acquire 得到全新条目
        let s3 = pool.slot("k");
        assert!(!Arc::ptr_eq(&s3, &s1));
        pool.release("k");
    }

    #[test]
    fn independent_keys_are_independent() {
        let pool = ConnectionPool::default();
        let a = pool.slot("a");
        let b = pool.slot("b");
        assert!(!Arc::ptr_eq(&a, &b));
        assert!(pool.release("a").is_some());
        assert!(pool.release("b").is_some());
    }

    #[test]
    fn ready_flag_and_hops_follow_entry_lifecycle() {
        let pool = ConnectionPool::default();
        assert!(!pool.is_ready("k"));
        pool.slot("k");
        pool.slot("k"); // 两个会话共享
        pool.mark_ready("k");
        assert!(pool.is_ready("k"));
        pool.attach_hops("k", Vec::new());
        // 释放一个引用后条目仍在（另一个会话还在用），ready 保持
        assert!(pool.release("k").is_none());
        assert!(pool.is_ready("k"));
        // 最后一个引用释放后条目消失
        assert!(pool.release("k").is_some());
        assert!(!pool.is_ready("k"));
    }

    #[test]
    fn over_release_is_safe() {
        let pool = ConnectionPool::default();
        assert!(pool.release("ghost").is_none());
        pool.slot("k");
        pool.release("k");
        // 多余的 release 不应 panic（saturating 语义）
        let _ = pool.release("k");
    }

    #[test]
    fn pool_key_distinguishes_identity() {
        let k1 = pool_key_of("ops", "srv", 22, "key", Some("keytext"));
        let k2 = pool_key_of("ops", "srv", 22, "key", Some("other-key"));
        let k3 = pool_key_of("ops", "srv", 22, "key", Some("keytext"));
        let k4 = pool_key_of("root", "srv", 22, "key", Some("keytext"));
        assert_ne!(k1, k2, "不同认证材料不应共享连接");
        assert_eq!(k1, k3, "相同身份应命中同一池条目");
        assert_ne!(k1, k4, "不同用户不应共享连接");
        // 无认证材料（理论上不出现）也不与他人混淆
        let k5 = pool_key_of("ops", "srv", 22, "password", None);
        assert_ne!(k1, k5);
    }
}