use std::fs;

use ssh2::{CheckResult, HostKeyType, KnownHostFileKind, KnownHostKeyFormat, Session};

pub(super) fn verify_host_key(session: &Session, host: &str, port: u16) -> Result<(), String> {
    let (key, key_type) = session
        .host_key()
        .ok_or_else(|| "服务器未提供主机密钥".to_string())?;
    let Some(home) = dirs::home_dir() else {
        return Err("无法确定用户主目录，不能校验 known_hosts".to_string());
    };
    let ssh_dir = home.join(".ssh");
    fs::create_dir_all(&ssh_dir).map_err(|error| format!("创建 ~/.ssh 失败: {error}"))?;
    let path = ssh_dir.join("known_hosts");
    let mut known = session
        .known_hosts()
        .map_err(|error| format!("初始化 known_hosts 失败: {error}"))?;
    if path.is_file() {
        known
            .read_file(&path, KnownHostFileKind::OpenSSH)
            .map_err(|error| format!("读取 known_hosts 失败: {error}"))?;
    }
    match known.check_port(host, port, key) {
        CheckResult::Match => Ok(()),
        CheckResult::Mismatch => Err(format!(
            "主机密钥已变化，已拒绝连接。请确认安全后从 known_hosts 删除 {}:{} 的旧记录",
            host, port
        )),
        CheckResult::NotFound | CheckResult::Failure => {
            let host_label = known_host_label(host, port);
            known
                .add(&host_label, key, &host_label, known_host_format(key_type))
                .map_err(|error| format!("记录主机密钥失败: {error}"))?;
            known
                .write_file(&path, KnownHostFileKind::OpenSSH)
                .map_err(|error| format!("写入 known_hosts 失败: {error}"))
        }
    }
}

fn known_host_label(host: &str, port: u16) -> String {
    if port == 22 {
        host.to_string()
    } else {
        format!("[{host}]:{port}")
    }
}

fn known_host_format(key_type: HostKeyType) -> KnownHostKeyFormat {
    match key_type {
        HostKeyType::Rsa => KnownHostKeyFormat::SshRsa,
        HostKeyType::Dss => KnownHostKeyFormat::SshDss,
        HostKeyType::Ecdsa256 => KnownHostKeyFormat::Ecdsa256,
        HostKeyType::Ecdsa384 => KnownHostKeyFormat::Ecdsa384,
        HostKeyType::Ecdsa521 => KnownHostKeyFormat::Ecdsa521,
        HostKeyType::Ed25519 => KnownHostKeyFormat::Ed25519,
        HostKeyType::Unknown => KnownHostKeyFormat::Unknown,
    }
}

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

    #[test]
    fn formats_non_default_ports_like_openssh() {
        assert_eq!(known_host_label("example.com", 22), "example.com");
        assert_eq!(known_host_label("example.com", 2222), "[example.com]:2222");
    }
}