// SSH密钥交换完整流程（Phase 3剩余）
// 参考OpenSSH kex.c: complete implementation

use crate::ssh_server::crypto::SessionKeys;
use crate::ssh_server::kex::KexResult;
use crate::ssh_server::kex_exchange::KexExchangeHandler;
use crate::ssh_server::packet::{PacketType, SshPacket};
use anyhow::{anyhow, Result};
use log::info;

/// SSH密钥交换完整状态管理（参考OpenSSH struct kex）
pub struct KexState {
    pub client_version: String,
    pub server_version: String,
    pub client_kexinit_payload: Vec<u8>,
    pub server_kexinit_payload: Vec<u8>,
    pub exchange_handler: KexExchangeHandler,
    pub session_keys: Option<SessionKeys>,
    pub newkeys_received: bool,
    pub newkeys_sent: bool,
}

impl KexState {
    /// 创建密钥交换状态
    pub fn new(
        client_version: String,
        server_version: String,
        kex_result: KexResult,
    ) -> Result<Self> {
        let exchange_handler = KexExchangeHandler::new(kex_result)?;

        Ok(Self {
            client_version,
            server_version,
            client_kexinit_payload: Vec::new(),
            server_kexinit_payload: Vec::new(),
            exchange_handler,
            session_keys: None,
            newkeys_received: false,
            newkeys_sent: false,
        })
    }

    /// 保存KEXINIT payloads（用于Exchange Hash计算）
    ///
    /// 分析OpenSSH源码后的结论：
    /// - kex->peer存储的是：incoming_packet剩余内容（payload fields + padding）
    /// - kex->my存储的是：prop2buf()结果（payload fields，不包括padding）
    ///
    /// **但exchange hash必须使用相同的I_C/I_S！**
    ///
    /// 疑问：OpenSSH如何确保client和server使用相同的padding？
    /// 可能答案：OpenSSH在计算exchange hash时，不包括padding？
    ///
    /// 暂时保持不包括padding（因为签名验证之前成功）
    pub fn save_kexinit_payloads(
        &mut self,
        client_kexinit: &SshPacket,
        server_kexinit: &SshPacket,
    ) {
        // Only save payload (without padding) for now
        self.client_kexinit_payload = client_kexinit.payload.clone();
        self.server_kexinit_payload = server_kexinit.payload.clone();

        info!("Saved KEXINIT payloads (payload only, no padding)");
        info!(
            "  client payload: {} bytes",
            self.client_kexinit_payload.len()
        );
        info!(
            "  server payload: {} bytes",
            self.server_kexinit_payload.len()
        );
    }

    /// 处理SSH_MSG_NEWKEYS（参考OpenSSH kex.c: kex_input_newkeys()）
    pub fn handle_newkeys(&mut self, packet: &SshPacket) -> Result<()> {
        info!("Processing SSH_MSG_NEWKEYS");

        // 验证packet类型
        if packet.payload.is_empty() {
            return Err(anyhow!("Invalid NEWKEYS packet"));
        }

        let packet_type = packet.payload[0];
        if packet_type != PacketType::SSH_MSG_NEWKEYS as u8 {
            return Err(anyhow!("Invalid packet type for NEWKEYS"));
        }

        // 标记NEWKEYS接收完成（参考OpenSSH）
        self.newkeys_received = true;

        info!("SSH_MSG_NEWKEYS received, encryption channel ready");

        Ok(())
    }

    /// 发送SSH_MSG_NEWKEYS（参考OpenSSH kex.c: kex_send_newkeys()）
    pub fn send_newkeys() -> Result<SshPacket> {
        info!("Sending SSH_MSG_NEWKEYS");

        let payload = vec![PacketType::SSH_MSG_NEWKEYS as u8];

        Ok(SshPacket::new(payload))
    }

    /// 检查NEWKEYS完成状态（加密通道建立）
    pub fn is_encryption_ready(&self) -> bool {
        self.newkeys_received && self.newkeys_sent
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ssh_server::kex::{KexProposal, KexResult};

    #[test]
    fn test_newkeys_handling() {
        let kex_result = KexResult::choose_algorithms(
            &KexProposal::server_default(),
            &KexProposal::client_default(),
        )
        .unwrap();

        let mut state = KexState::new(
            "SSH-2.0-OpenSSH_10.2".to_string(),
            "SSH-2.0-MarkBaseSSH_1.0".to_string(),
            kex_result,
        )
        .unwrap();

        let newkeys_packet = SshPacket::new(vec![PacketType::SSH_MSG_NEWKEYS as u8]);

        state.handle_newkeys(&newkeys_packet).unwrap();

        assert!(state.newkeys_received);
    }
}