// SSH密钥交换算法协商实现（Phase 2）
// 参考OpenSSH kex.c: kex_send_kexinit(), kex_choose_conf()

use crate::ssh_server::packet::{SshPacket, PacketType};
use anyhow::{Result, anyhow};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use log::{info, debug};
use std::io::{Read, Write};

/// SSH算法类型（参考OpenSSH PROTOCOL定义）
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum AlgorithmType {
    KEX_ALGS = 0,          // 密钥交换算法
    SERVER_HOST_KEY_ALGS = 1, // 服务器主机密钥算法
    ENC_ALGS_CTOS = 2,     // 客户端到服务器加密算法
    ENC_ALGS_STOC = 3,     // 服务器到客户端加密算法
    MAC_ALGS_CTOS = 4,     // 客户端到服务器MAC算法
    MAC_ALGS_STOC = 5,     // 服务器到客户端MAC算法
    COMP_ALGS_CTOS = 6,    // 客户端到服务器压缩算法
    COMP_ALGS_STOC = 7,    // 服务器到客户端压缩算法
    LANGS_CTOS = 8,        // 客户端到服务器语言
    LANGS_STOC = 9,        // 服务器到客户端语言
}

/// SSH算法提议（参考OpenSSH kex.h: struct kex）
#[derive(Debug, Clone)]
pub struct KexProposal {
    pub kex_algorithms: String,         // 密钥交换算法列表
    pub server_host_key_algorithms: String, // 主机密钥算法列表
    pub encryption_algorithms_ctos: String, // 加密算法（客户端→服务器）
    pub encryption_algorithms_stoc: String, // 加密算法（服务器→客户端）
    pub mac_algorithms_ctos: String,    // MAC算法（客户端→服务器）
    pub mac_algorithms_stoc: String,    // MAC算法（服务器→客户端）
    pub compression_algorithms_ctos: String, // 压缩算法（客户端→服务器）
    pub compression_algorithms_stoc: String, // 压缩算法（服务器→客户端）
    pub languages_ctos: String,         // 语言（客户端→服务器）
    pub languages_stoc: String,         // 语言（服务器→客户端）
    pub first_kex_packet_follows: bool, // 是否立即发送第一个KEX packet
    pub reserved: u32,                  // 保留字段（0）
}

impl KexProposal {
    /// 创建默认算法提议（参考OpenSSH myproposal.h）
    pub fn server_default() -> Self {
        // 参考OpenSSH KEX_SERVER定义
        Self {
            // 密钥交换算法：优先Curve25519（推荐）
            kex_algorithms: "curve25519-sha256,curve25519-sha256@libssh.org,diffie-hellman-group14-sha256".to_string(),
            
            // 主机密钥算法：优先Ed25519
            server_host_key_algorithms: "ssh-ed25519,rsa-sha2-256,rsa-sha2-512".to_string(),
            
            // 加密算法：AES-256-CTR（推荐）
            encryption_algorithms_ctos: "aes256-ctr,aes128-ctr".to_string(),
            encryption_algorithms_stoc: "aes256-ctr,aes128-ctr".to_string(),
            
            // MAC算法：HMAC-SHA256
            mac_algorithms_ctos: "hmac-sha2-256,hmac-sha2-512".to_string(),
            mac_algorithms_stoc: "hmac-sha2-256,hmac-sha2-512".to_string(),
            
            // 压缩算法：none优先
            compression_algorithms_ctos: "none,zlib".to_string(),
            compression_algorithms_stoc: "none,zlib".to_string(),
            
            // 语言：空
            languages_ctos: "".to_string(),
            languages_stoc: "".to_string(),
            
            first_kex_packet_follows: false,
            reserved: 0,
        }
    }
    
    /// 创建客户端默认提议（用于测试）
    pub fn client_default() -> Self {
        Self {
            kex_algorithms: "curve25519-sha256,diffie-hellman-group14-sha256".to_string(),
            server_host_key_algorithms: "ssh-ed25519,rsa-sha2-256".to_string(),
            encryption_algorithms_ctos: "aes256-ctr,aes128-ctr".to_string(),
            encryption_algorithms_stoc: "aes256-ctr,aes128-ctr".to_string(),
            mac_algorithms_ctos: "hmac-sha2-256".to_string(),
            mac_algorithms_stoc: "hmac-sha2-256".to_string(),
            compression_algorithms_ctos: "none".to_string(),
            compression_algorithms_stoc: "none".to_string(),
            languages_ctos: "".to_string(),
            languages_stoc: "".to_string(),
            first_kex_packet_follows: false,
            reserved: 0,
        }
    }
    
    /// 序列化到SSH_MSG_KEXINIT packet（参考OpenSSH kex_send_kexinit()）
    pub fn to_kexinit_packet(&self) -> Result<SshPacket> {
        let mut payload = Vec::new();
        
        // Packet type
        payload.write_u8(PacketType::SSH_MSG_KEXINIT as u8)?;
        
        // Cookie（16字节随机数，OpenSSH要求）
        // 简化：使用固定值（实际应随机生成）
        let cookie = [0u8; 16];
        payload.write_all(&cookie)?;
        
        // 10个算法列表（SSH string格式：length + data）
        write_ssh_string(&mut payload, &self.kex_algorithms)?;
        write_ssh_string(&mut payload, &self.server_host_key_algorithms)?;
        write_ssh_string(&mut payload, &self.encryption_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.encryption_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.mac_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.mac_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.compression_algorithms_ctos)?;
        write_ssh_string(&mut payload, &self.compression_algorithms_stoc)?;
        write_ssh_string(&mut payload, &self.languages_ctos)?;
        write_ssh_string(&mut payload, &self.languages_stoc)?;
        
        // first_kex_packet_follows（boolean）
        payload.write_u8(if self.first_kex_packet_follows { 1 } else { 0 })?;
        
        // reserved（u32）
        payload.write_u32::<BigEndian>(self.reserved)?;
        
        Ok(SshPacket::new(payload))
    }
    
    /// 从SSH_MSG_KEXINIT packet解析（参考OpenSSH kex_input_kexinit()）
    pub fn from_kexinit_packet(packet: &SshPacket) -> Result<Self> {
        let mut cursor = std::io::Cursor::new(packet.payload.as_slice());  // 使用as_slice()（Rust标准）
        
        // Packet type
        let packet_type = cursor.read_u8()?;
        if packet_type != PacketType::SSH_MSG_KEXINIT as u8 {
            return Err(anyhow!("Invalid packet type for KEXINIT"));
        }
        
        // Cookie（16字节，忽略）
        cursor.read_exact(&mut [0u8; 16])?;
        
        // 10个算法列表
        let kex_algorithms = read_ssh_string(&mut cursor)?;
        let server_host_key_algorithms = read_ssh_string(&mut cursor)?;
        let encryption_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let encryption_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let mac_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let mac_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let compression_algorithms_ctos = read_ssh_string(&mut cursor)?;
        let compression_algorithms_stoc = read_ssh_string(&mut cursor)?;
        let languages_ctos = read_ssh_string(&mut cursor)?;
        let languages_stoc = read_ssh_string(&mut cursor)?;
        
        // first_kex_packet_follows
        let first_kex_packet_follows = cursor.read_u8()? != 0;
        
        // reserved
        let reserved = cursor.read_u32::<BigEndian>()?;
        
        Ok(Self {
            kex_algorithms,
            server_host_key_algorithms,
            encryption_algorithms_ctos,
            encryption_algorithms_stoc,
            mac_algorithms_ctos,
            mac_algorithms_stoc,
            compression_algorithms_ctos,
            compression_algorithms_stoc,
            languages_ctos,
            languages_stoc,
            first_kex_packet_follows,
            reserved,
        })
    }
}

/// SSH算法协商结果（参考OpenSSH struct kex）
#[derive(Debug, Clone)]
pub struct KexResult {
    pub kex_algorithm: String,          // 选定的密钥交换算法
    pub host_key_algorithm: String,     // 选定的主机密钥算法
    pub encryption_ctos: String,        // 选定的加密算法（客户端→服务器）
    pub encryption_stoc: String,        // 选定的加密算法（服务器→客户端）
    pub mac_ctos: String,               // 选定的MAC算法（客户端→服务器）
    pub mac_stoc: String,               // 选定的MAC算法（服务器→客户端）
    pub compression_ctos: String,       // 选定的压缩算法（客户端→服务器）
    pub compression_stoc: String,       // 选定的压缩算法（服务器→客户端）
}

/// 算法匹配逻辑（参考OpenSSH kex_choose_conf()）
impl KexResult {
    /// 从服务器和客户端提议中选择算法（参考OpenSSH kex_choose_conf()）
    pub fn choose_algorithms(server: &KexProposal, client: &KexProposal) -> Result<Self> {
        info!("Starting algorithm negotiation");
        
        // 算法匹配：优先客户端偏好（OpenSSH逻辑）
        // 参考OpenSSH：客户端列出的算法顺序为偏好顺序
        
        // 密钥交换算法匹配
        let kex_algorithm = match_algorithm(&client.kex_algorithms, &server.kex_algorithms)?;
        
        // 主机密钥算法匹配
        let host_key_algorithm = match_algorithm(&client.server_host_key_algorithms, &server.server_host_key_algorithms)?;
        
        // 加密算法匹配
        let encryption_ctos = match_algorithm(&client.encryption_algorithms_ctos, &server.encryption_algorithms_ctos)?;
        let encryption_stoc = match_algorithm(&client.encryption_algorithms_stoc, &server.encryption_algorithms_stoc)?;
        
        // MAC算法匹配
        let mac_ctos = match_algorithm(&client.mac_algorithms_ctos, &server.mac_algorithms_ctos)?;
        let mac_stoc = match_algorithm(&client.mac_algorithms_stoc, &server.mac_algorithms_stoc)?;
        
        // 压缩算法匹配
        let compression_ctos = match_algorithm(&client.compression_algorithms_ctos, &server.compression_algorithms_ctos)?;
        let compression_stoc = match_algorithm(&client.compression_algorithms_stoc, &server.compression_algorithms_stoc)?;
        
        info!("Algorithm negotiation completed:");
        debug!("  KEX: {}", kex_algorithm);
        debug!("  Host key: {}", host_key_algorithm);
        debug!("  Encryption (C->S): {}", encryption_ctos);
        debug!("  Encryption (S->C): {}", encryption_stoc);
        debug!("  MAC (C->S): {}", mac_ctos);
        debug!("  MAC (S->C): {}", mac_stoc);
        
        Ok(Self {
            kex_algorithm,
            host_key_algorithm,
            encryption_ctos,
            encryption_stoc,
            mac_ctos,
            mac_stoc,
            compression_ctos,
            compression_stoc,
        })
    }
}

/// 算法匹配函数（参考OpenSSH match.c: match_list()）
fn match_algorithm(client_algs: &str, server_algs: &str) -> Result<String> {
    // 算法列表格式：name1,name2,name3,...
    let client_list: Vec<&str> = client_algs.split(',').collect();
    let server_list: Vec<&str> = server_algs.split(',').collect();
    
    // OpenSSH逻辑：按客户端偏好顺序匹配
    for client_alg in &client_list {
        if server_list.contains(client_alg) {
            return Ok(client_alg.to_string());
        }
    }
    
    Err(anyhow!("No matching algorithm found: client={}, server={}", client_algs, server_algs))
}

/// SSH string写入辅助函数（length + data）
fn write_ssh_string<W: Write>(writer: &mut W, s: &str) -> Result<()> {
    writer.write_u32::<BigEndian>(s.len() as u32)?;
    writer.write_all(s.as_bytes())?;
    Ok(())
}

/// SSH string读取辅助函数（length + data）
fn read_ssh_string<R: Read>(reader: &mut R) -> Result<String> {
    let length = reader.read_u32::<BigEndian>()?;
    let mut buffer = vec![0u8; length as usize];
    reader.read_exact(&mut buffer)?;
    Ok(String::from_utf8(buffer)?)
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_kex_proposal_creation() {
        let proposal = KexProposal::server_default();
        assert!(proposal.kex_algorithms.contains("curve25519-sha256"));
    }
    
    #[test]
    fn test_kex_proposal_serialization() {
        let proposal = KexProposal::server_default();
        let packet = proposal.to_kexinit_packet().unwrap();
        assert!(packet.payload.len() > 0);
    }
    
    #[test]
    fn test_algorithm_matching() {
        let client = "curve25519-sha256,aes256-ctr";
        let server = "aes256-ctr,diffie-hellman-group14-sha256";
        
        let matched = match_algorithm(client, server).unwrap();
        assert_eq!(matched, "aes256-ctr");  // 按客户端顺序匹配
    }
    
    #[test]
    fn test_kex_negotiation() {
        let server = KexProposal::server_default();
        let client = KexProposal::client_default();
        
        let result = KexResult::choose_algorithms(&server, &client).unwrap();
        assert_eq!(result.kex_algorithm, "curve25519-sha256");  // 优先Curve25519
        assert_eq!(result.encryption_ctos, "aes256-ctr");  // AES-256-CTR
    }
}
