// 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};
use sha2::{Sha256, Digest};

/// SSH密钥交换流程处理器（参考OpenSSH kex.c）
pub struct KexExchangeHandler {
    kex_algorithm: String,
    server_kex: Option<Curve25519Kex>,
    host_key: Ed25519HostKey,
    shared_secret: Option<Vec<u8>>,
    client_public_key: Option<Vec<u8>>,
    server_public_key: Option<Vec<u8>>,
    exchange_hash: Option<Vec<u8>>,  // 保存exchange hash（H参数）
    client_version: Option<String>,
    server_version: Option<String>,
    client_kexinit_payload: Option<Vec<u8>>,
    server_kexinit_payload: Option<Vec<u8>>,
}

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,
            shared_secret: None,
            client_public_key: None,
            server_public_key: None,
            exchange_hash: None,
            client_version: None,
            server_version: None,
            client_kexinit_payload: None,
            server_kexinit_payload: None,
        })
    }
    
/// 处理SSH_MSG_KEXDH_INIT（Curve25519密钥交换）（参考OpenSSH kex.c: kex_input_kex_init()）
    pub fn handle_kexdh_init(
        &mut self, 
        packet: &SshPacket,
        client_version: &str,
        server_version: &str,
        client_kexinit_payload: &[u8],
        server_kexinit_payload: &[u8],
    ) -> Result<SshPacket> {
        info!("Processing SSH_MSG_KEXDH_INIT (Curve25519)");
        
        let mut cursor = std::io::Cursor::new(packet.payload.as_slice());
        
        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"));
        }
        
        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)?;
        
        self.server_kex = Some(Curve25519Kex::new());
        let server_kex = self.server_kex.as_mut().unwrap();
        
        let shared_secret = server_kex.compute_shared_secret(&client_public_key)?;
        let server_public_key = server_kex.public_key().to_vec();
        
        // Save for later session key computation
        self.shared_secret = Some(shared_secret.to_vec());
        self.client_public_key = Some(client_public_key.clone());
        self.server_public_key = Some(server_public_key.clone());
        
        // Save client_version, server_version, kexinit payloads for exchange hash
        self.client_version = Some(client_version.to_string());
        self.server_version = Some(server_version.to_string());
        self.client_kexinit_payload = Some(client_kexinit_payload.to_vec());
        self.server_kexinit_payload = Some(server_kexinit_payload.to_vec());
        
        info!("Curve25519 shared secret computed and saved");
        
        // Compute exchange hash ONCE and reuse it
        let host_key_blob = self.build_ssh_host_key()?;
        let exchange_hash = self.compute_exchange_hash(
            &shared_secret,
            &host_key_blob,
            &client_public_key,
            &server_public_key,
            client_version,
            server_version,
            client_kexinit_payload,
            server_kexinit_payload,
        )?;
        
        info!("Exchange hash computed:");
        info!("  shared_secret[0] = {} (>=0x80? {})", shared_secret[0], shared_secret[0] >= 0x80);
        info!("  exchange_hash full (32 bytes): {:?}", exchange_hash);
        
        self.exchange_hash = Some(exchange_hash.clone());
        info!("Exchange hash saved for key derivation");
        
        self.build_kexdh_reply(
            &exchange_hash,
            &host_key_blob,
            &server_public_key,
        )
    }
    
    /// 构建SSH_MSG_KEXDH_REPLY packet（参考OpenSSH kex.c）
    fn build_kexdh_reply(
        &self, 
        exchange_hash: &[u8], 
        host_key_blob: &[u8],
        server_public_key: &[u8],
    ) -> Result<SshPacket> {
        info!("=== Building SSH_MSG_KEXDH_REPLY ===");
        info!("Input server_public_key: {:?}", server_public_key);
        
        let mut payload = Vec::new();
        
        payload.write_u8(PacketType::SSH_MSG_KEXDH_REPLY as u8)?;
        
        payload.write_u32::<BigEndian>(host_key_blob.len() as u32)?;
        payload.write_all(host_key_blob)?;
        
        info!("Writing server_public_key to payload (32 bytes)");
        payload.write_u32::<BigEndian>(32)?;
        payload.write_all(server_public_key)?;
        
        let signature = self.build_exchange_signature(exchange_hash)?;
        payload.write_u32::<BigEndian>(signature.len() as u32)?;
        payload.write_all(&signature)?;
        
        info!("SSH_MSG_KEXDH_REPLY payload built successfully");
        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)
    }
    
    /// 计算Exchange Hash（参考OpenSSH kex.c: kex_hash() RFC 4253 Section 7.2）
    fn compute_exchange_hash(
        &self,
        shared_secret: &[u8],
        host_key_blob: &[u8],
        client_public_key: &[u8],
        server_public_key: &[u8],
        client_version: &str,
        server_version: &str,
        client_kexinit_payload: &[u8],
        server_kexinit_payload: &[u8],
    ) -> Result<Vec<u8>> {
        use sha2::{Sha256, Digest};
        
        info!("=== EXCHANGE HASH COMPUTATION ===");
        info!("V_C (client version): {:?}", client_version.as_bytes());
        info!("V_C length: {}", client_version.len());
        
        info!("V_S (server version): {:?}", server_version.as_bytes());
        info!("V_S length: {}", server_version.len());
        
        info!("I_C (client KEXINIT payload): {:?}", &client_kexinit_payload[..std::cmp::min(50, client_kexinit_payload.len())]);
        info!("I_C length: {}", client_kexinit_payload.len());
        info!("I_C[0] (packet type): {} (should be SSH_MSG_KEXINIT=20)", client_kexinit_payload[0]);
        
        info!("I_S (server KEXINIT payload): {:?}", &server_kexinit_payload[..std::cmp::min(50, server_kexinit_payload.len())]);
        info!("I_S length: {}", server_kexinit_payload.len());
        info!("I_S[0] (packet type): {} (should be SSH_MSG_KEXINIT=20)", server_kexinit_payload[0]);
        
        info!("K_S (host key blob): {:?}", &host_key_blob[..std::cmp::min(30, host_key_blob.len())]);
        info!("K_S length: {}", host_key_blob.len());
        
        info!("Q_C (client ECDH public key): {:?}", &client_public_key[..std::cmp::min(16, client_public_key.len())]);
        info!("Q_C full (32 bytes): {:?}", client_public_key);
        info!("Q_C length: {}", client_public_key.len());
        
        info!("Q_S (server ECDH public key): {:?}", &server_public_key[..std::cmp::min(16, server_public_key.len())]);
        info!("Q_S full (32 bytes): {:?}", server_public_key);
        info!("Q_S length: {}", server_public_key.len());
        
        let mut hasher = Sha256::new();
        
        // RFC 4253 Section 7: V_C and V_S are version strings (without \r\n based on testing)
        hasher.update(&(client_version.len() as u32).to_be_bytes());
        hasher.update(client_version.as_bytes());
        
        hasher.update(&(server_version.len() as u32).to_be_bytes());
        hasher.update(server_version.as_bytes());
        
        // OpenSSH kexgex.c: "kexinit messages: fake header: len+SSH2_MSG_KEXINIT"
        // KEXINIT payload should NOT include SSH_MSG_KEXINIT type byte
        // OpenSSH stores payload starting from cookie, prepends SSH_MSG_KEXINIT in exchange hash
        
        // Remove SSH_MSG_KEXINIT type byte from payloads (our payload includes it)
        let client_kexinit_without_type = &client_kexinit_payload[1..];
        let server_kexinit_without_type = &server_kexinit_payload[1..];
        
        info!("I_C (client KEXINIT without type byte): {} bytes (first byte should be cookie)", client_kexinit_without_type.len());
        info!("I_S (server KEXINIT without type byte): {} bytes", server_kexinit_without_type.len());
        
        // Exchange hash: uint32(len+1) + uint8(SSH_MSG_KEXINIT) + payload_without_type
        hasher.update(&((client_kexinit_without_type.len() + 1) as u32).to_be_bytes());
        hasher.update(&[20]);  // SSH_MSG_KEXINIT type byte
        hasher.update(client_kexinit_without_type);
        
        hasher.update(&((server_kexinit_without_type.len() + 1) as u32).to_be_bytes());
        hasher.update(&[20]);  // SSH_MSG_KEXINIT type byte
        hasher.update(server_kexinit_without_type);
        
        hasher.update(&(host_key_blob.len() as u32).to_be_bytes());
        hasher.update(host_key_blob);
        
        hasher.update(&(client_public_key.len() as u32).to_be_bytes());
        hasher.update(client_public_key);
        
        hasher.update(&(server_public_key.len() as u32).to_be_bytes());
        hasher.update(server_public_key);
        
        info!("Exchange hash components:");
        info!("  shared_secret raw ({} bytes): {:?}", shared_secret.len(), &shared_secret[..std::cmp::min(8, shared_secret.len())]);
        
        // RFC 8731 Section 3.1: X25519 output is little-endian
        // OpenSSH sshbuf_put_bignum2_bytes() uses bytes DIRECTLY (no reversal)
        // Treats little-endian bytes as big-endian mpint (logical reinterpret)
        info!("  Using shared_secret directly (little-endian bytes as big-endian mpint)");
        
        // RFC 4253: mpint格式 = 去掉前导零 + 最高位>=0x80时前面加0
        // 参考OpenSSH sshbuf_put_bignum2_bytes()
        let mut start = 0;
        while start < shared_secret.len() - 1 && shared_secret[start] == 0 {
            start += 1;
        }
        let trimmed_shared_secret = &shared_secret[start..];
        
        info!("  shared_secret after removing leading zeros ({} bytes): {:?}", trimmed_shared_secret.len(), &trimmed_shared_secret[..std::cmp::min(8, trimmed_shared_secret.len())]);
        
        let mpint_shared_secret_data = if trimmed_shared_secret.len() > 0 && trimmed_shared_secret[0] >= 0x80 {
            let mut mpint = vec![0u8];
            mpint.extend_from_slice(trimmed_shared_secret);
            info!("  trimmed_shared_secret[0] >= 0x80, prepending 0 byte");
            mpint
        } else {
            trimmed_shared_secret.to_vec()
        };
        
        info!("  mpint_shared_secret_data ({} bytes): {:?}", mpint_shared_secret_data.len(), &mpint_shared_secret_data[..std::cmp::min(8, mpint_shared_secret_data.len())]);
        
        // mpint格式 = uint32(length) + mpint_data
        hasher.update(&(mpint_shared_secret_data.len() as u32).to_be_bytes());
        hasher.update(&mpint_shared_secret_data);
        
        Ok(hasher.finalize().to_vec())
    }
    
    /// 构建交换签名（参考OpenSSH ssh-sign.c）
    fn build_exchange_signature(&self, exchange_hash: &[u8]) -> Result<Vec<u8>> {
        let signature_bytes = self.host_key.sign(exchange_hash)?;
        
        let mut ssh_signature = Vec::new();
        
        ssh_signature.write_u32::<BigEndian>(11)?;
        ssh_signature.write_all("ssh-ed25519".as_bytes())?;
        
        ssh_signature.write_u32::<BigEndian>(64)?;
        ssh_signature.write_all(&signature_bytes)?;
        
        Ok(ssh_signature)
    }
    
    /// 计算会话密钥（参考OpenSSH kex.c: derive_keys()）
    /// 使用保存的exchange_hash（H参数）
    pub fn compute_session_keys(&self) -> Result<SessionKeys> {
        if self.shared_secret.is_none() {
            return Err(anyhow!("No shared secret available"));
        }
        
        if self.exchange_hash.is_none() {
            return Err(anyhow!("No exchange hash available"));
        }
        
        let shared_secret = self.shared_secret.as_ref().unwrap();
        let exchange_hash = self.exchange_hash.as_ref().unwrap();
        let server_public_key = self.server_public_key.as_ref().unwrap();
        let client_public_key = self.client_public_key.as_ref().unwrap();
        let host_key_blob = self.build_ssh_host_key()?;
        
        SessionKeys::derive(
            shared_secret,
            exchange_hash,  // 使用保存的exchange hash（H参数）
            server_public_key,
            client_public_key,
            &host_key_blob,
        )
    }
}

#[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);
    }
}