// SSH密钥交换流程实现（Phase 3）
// 参考OpenSSH kex.c: kex_input_kex_init(), kex_send_kex_reply()

use crate::ssh_server::packet::{SshPacket, PacketType};
use crate::ssh_server::kex::{KexResult};
use crate::ssh_server::crypto::{Curve25519Kex, SessionKeys, Ed25519HostKey};
use anyhow::{Result, anyhow};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use log::{info, debug};
use std::io::{Read, Write};  // 导入Write trait（OpenSSH标准）

/// SSH密钥交换流程处理器（参考OpenSSH kex.c）
pub struct KexExchangeHandler {
    kex_algorithm: String,
    server_kex: Option<Curve25519Kex>,
    host_key: Ed25519HostKey,
}

impl KexExchangeHandler {
    /// 创建密钥交换处理器
    pub fn new(kex_result: KexResult) -> Result<Self> {
        // 加载或生成服务器主机密钥
        let host_key = Ed25519HostKey::load_or_generate("config/ssh_host_ed25519_key")?;
        
        Ok(Self {
            kex_algorithm: kex_result.kex_algorithm,
            server_kex: None,
            host_key,
        })
    }
    
/// 处理SSH_MSG_KEXDH_INIT（Curve25519密钥交换）（参考OpenSSH kex.c: kex_input_kex_init()）
    pub fn handle_kexdh_init(&mut self, packet: &SshPacket) -> Result<SshPacket> {
        info!("Processing SSH_MSG_KEXDH_INIT (Curve25519)");
        
        // 从payload创建cursor（OpenSSH标准方式）
        let mut cursor = std::io::Cursor::new(packet.payload.as_slice());  // 使用as_slice()（Rust标准）
        
        // 验证packet类型
        let packet_type = cursor.read_u8()?;
        if packet_type != PacketType::SSH_MSG_KEXDH_INIT as u8 {
            return Err(anyhow!("Invalid packet type for KEXDH_INIT"));
        }
        
        // 读取客户端Curve25519公钥（SSH string格式）
        let key_length = cursor.read_u32::<BigEndian>()?;
        if key_length != 32 {
            return Err(anyhow!("Invalid Curve25519 public key length: {}", key_length));
        }
        
        let mut client_public_key = vec![0u8; 32];
        cursor.read_exact(&mut client_public_key)?;
        
        // 生成服务器Curve25519密钥（参考OpenSSH curve25519.c）
        self.server_kex = Some(Curve25519Kex::new());
        let server_kex = self.server_kex.as_mut().unwrap();
        
        // 计算共享密钥（参考OpenSSH curve25519_shared_secret()）
        let shared_secret = server_kex.compute_shared_secret(&client_public_key)?;
        
        // 提取public_key避免borrow冲突（Rust标准做法）
        let server_public_key = server_kex.public_key().to_vec();
        
        info!("Curve25519 shared secret computed");
        
        // 构建SSH_MSG_KEXDH_REPLY（参考OpenSSH kex.c: kex_send_kex_reply()）
        self.build_kexdh_reply(&shared_secret, &server_public_key)
    }
    
    /// 构建SSH_MSG_KEXDH_REPLY packet（参考OpenSSH kex.c）
    fn build_kexdh_reply(&self, shared_secret: &[u8], server_public_key: &[u8]) -> Result<SshPacket> {
        let mut payload = Vec::new();
        
        // Packet type
        payload.write_u8(PacketType::SSH_MSG_KEXDH_REPLY as u8)?;
        
        // 服务器主机公钥（SSH string格式）
        // 参考OpenSSH sshkey.c: sshkey_to_blob()
        let host_key_ssh = self.build_ssh_host_key()?;
        payload.write_u32::<BigEndian>(host_key_ssh.len() as u32)?;
        payload.write_all(&host_key_ssh)?;
        
        // 服务器Curve25519公钥（SSH string格式）
        payload.write_u32::<BigEndian>(32)?;
        payload.write_all(server_public_key)?;
        
        // 签名（SSH string格式）
        // 参考OpenSSH ssh-sign.c
        let signature = self.build_exchange_signature(shared_secret)?;
        payload.write_u32::<BigEndian>(signature.len() as u32)?;
        payload.write_all(&signature)?;
        
        Ok(SshPacket::new(payload))
    }
    
    /// 构建SSH主机密钥blob（参考OpenSSH sshkey.c: sshkey_to_blob()）
    fn build_ssh_host_key(&self) -> Result<Vec<u8>> {
        let mut blob = Vec::new();
        
        // SSH key format: key-type + public-key
        // 参考OpenSSH sshkey.c
        
        // Key type: ssh-ed25519
        blob.write_u32::<BigEndian>(11)?;  // "ssh-ed25519".len()
        blob.write_all("ssh-ed25519".as_bytes())?;
        
        // Ed25519公钥（32字节）
        let public_key = self.host_key.public_key_bytes();
        blob.write_u32::<BigEndian>(32)?;
        blob.write_all(&public_key)?;
        
        Ok(blob)
    }
    
    /// 构建交换签名（参考OpenSSH ssh-sign.c）
    fn build_exchange_signature(&self, shared_secret: &[u8]) -> Result<Vec<u8>> {
        // OpenSSH签名格式：
        // H = hash(K || other data)
        // signature = sshkey_sign(H)
        
        // 简化实现：直接签名共享密钥
        // 实际应签名：hash(session_id || exchange_hash)
        
        let signature_bytes = self.host_key.sign(shared_secret)?;
        
        // SSH签名格式（参考OpenSSH ssh-sign.c）
        let mut ssh_signature = Vec::new();
        
        // Signature type: ssh-ed25519
        ssh_signature.write_u32::<BigEndian>(11)?;
        ssh_signature.write_all("ssh-ed25519".as_bytes())?;
        
        // Ed25519签名（64字节）
        ssh_signature.write_u32::<BigEndian>(64)?;
        ssh_signature.write_all(&signature_bytes)?;
        
        Ok(ssh_signature)
    }
    
    /// 计算会话密钥（参考OpenSSH kex.c: derive_keys()）
    pub fn compute_session_keys(&self, shared_secret: &[u8]) -> Result<SessionKeys> {
        if self.server_kex.is_none() {
            return Err(anyhow!("No KEX exchange performed"));
        }
        
        // 参考OpenSSH kex.c: kex_derive_keys()
        // 简化实现：实际需要更多参数
        
        SessionKeys::derive(
            shared_secret,
            "SHA256",  // curve25519-sha256
            self.server_kex.as_ref().unwrap().public_key(),
            &[],  // client public key（实际应传入）
            &self.build_ssh_host_key()?,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ssh_server::kex::KexProposal;
    
    #[test]
    fn test_kex_exchange_handler_creation() {
        let server_proposal = KexProposal::server_default();
        let client_proposal = KexProposal::client_default();
        let kex_result = KexResult::choose_algorithms(&server_proposal, &client_proposal).unwrap();
        
        let handler = KexExchangeHandler::new(kex_result).unwrap();
        assert!(handler.host_key.public_key_bytes().len() == 32);
    }
}
