// SSH加密模块（Phase 3：密钥交换）
// 参考OpenSSH curve25519.c, kex.c

use anyhow::{Result, anyhow};
use x25519_dalek::{EphemeralSecret, PublicKey, SharedSecret};
use ed25519_dalek::{SigningKey, VerifyingKey, Signature, Signer};
use sha2::{Sha256, Digest};
use log::{info, debug};
use rand::rngs::OsRng;

/// Curve25519密钥交换处理器（参考OpenSSH curve25519.c）
pub struct Curve25519Kex {
    secret: Option<EphemeralSecret>,  // 使用Option包装（一次性使用类型）
    public: PublicKey,
}

impl Curve25519Kex {
    /// 创建新的Curve25519密钥交换实例
    pub fn new() -> Self {
        // 参考OpenSSH curve25519.c: curve25519_make_key()
        // x25519-dalek 2.0标准API：使用random_from_rng
        let secret = EphemeralSecret::random_from_rng(OsRng);
        let public = PublicKey::from(&secret);
        
        Self { secret: Some(secret), public }  // Some包装
    }
    
    /// 获取公钥（用于SSH_MSG_KEX_ECDH_INIT）
    pub fn public_key(&self) -> &[u8] {
        self.public.as_bytes()
    }
    
    /// 计算共享密钥（参考OpenSSH curve25519_shared_secret()）
    /// 使用&mut self（消耗模式，符合OpenSSH设计）
    pub fn compute_shared_secret(&mut self, client_public: &[u8]) -> Result<[u8; 32]> {
        if client_public.len() != 32 {
            return Err(anyhow!("Invalid client public key length"));
        }
        
        // 参考OpenSSH：curve25519共享密钥计算
        let client_public = PublicKey::from(<[u8; 32]>::try_from(client_public)?);
        
        // 使用take()取出secret（Rust标准模式）
        if let Some(secret) = self.secret.take() {
            let shared_secret = secret.diffie_hellman(&client_public);
            Ok(shared_secret.as_bytes().clone())
        } else {
            Err(anyhow!("Secret already used"))
        }
    }
}

/// SSH会话密钥计算（参考OpenSSH kex.c: derive_keys()）
pub struct SessionKeys {
    pub session_id: Vec<u8>,
    pub encryption_key_ctos: Vec<u8>,
    pub encryption_key_stoc: Vec<u8>,
    pub mac_key_ctos: Vec<u8>,
    pub mac_key_stoc: Vec<u8>,
    pub iv_ctos: Vec<u8>,
    pub iv_stoc: Vec<u8>,
}

impl SessionKeys {
    /// 计算会话密钥（参考OpenSSH kex.c: kex_derive_keys()）
    /// RFC 4253 Section 7.2: Key = HASH(K || H || X || session_id)
    pub fn derive(
        shared_secret: &[u8],
        exchange_hash: &[u8],  // H参数（exchange hash）
        server_public_key: &[u8],
        client_public_key: &[u8],
        server_host_key: &[u8],
    ) -> Result<Self> {
        // RFC 4253: session_id = H (第一次exchange hash)
        let session_id = exchange_hash.to_vec();
        
        info!("SessionKeys::derive() starting");
        info!("  shared_secret ({} bytes): {:?}", shared_secret.len(), &shared_secret[..8]);
        info!("  shared_secret[0] = {} (>=0x80? {})", shared_secret[0], shared_secret[0] >= 0x80);
        
        // RFC 4253密钥派生公式：HASH(K || H || X || session_id)
        // 其中K是shared_secret（需要mpint格式）
        let shared_secret_mpint = Self::encode_mpint(shared_secret);
        
        info!("  shared_secret_mpint ({} bytes): {:?}", shared_secret_mpint.len(), &shared_secret_mpint[..std::cmp::min(12, shared_secret_mpint.len())]);
        
        let encryption_key_ctos = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'C', &session_id)?;
        let encryption_key_stoc = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'D', &session_id)?;
        let mac_key_ctos = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'E', &session_id)?;
        let mac_key_stoc = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'F', &session_id)?;
        
        let iv_ctos = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'A', &session_id)?;
        let iv_stoc = Self::derive_key_rfc4253(&shared_secret_mpint, exchange_hash, 'B', &session_id)?;
        
        Ok(Self {
            session_id,
            encryption_key_ctos,
            encryption_key_stoc,
            mac_key_ctos,
            mac_key_stoc,
            iv_ctos,
            iv_stoc,
        })
    }
    
    /// RFC 4253密钥派生函数
    /// 公式：Key = HASH(K || H || X || session_id)
    fn derive_key_rfc4253(K_mpint: &[u8], H: &[u8], X: char, session_id: &[u8]) -> Result<Vec<u8>> {
        let mut hasher = Sha256::new();
        
        info!("Deriving key for X='{}'", X);
        info!("  K_mpint ({} bytes): {:?}", K_mpint.len(), &K_mpint[..std::cmp::min(8, K_mpint.len())]);
        info!("  H ({} bytes): {:?}", H.len(), &H[..8]);
        info!("  session_id ({} bytes): {:?}", session_id.len(), &session_id[..8]);
        
        // RFC 4253: HASH(K || H || X || session_id)
        hasher.update(K_mpint);  // K (shared secret in mpint format)
        hasher.update(H);        // H (exchange hash)
        hasher.update(&[X as u8]); // X (single character)
        hasher.update(session_id); // session_id
        
        let full_hash = hasher.finalize();
        
        info!("  Derived key (first 8 bytes): {:?}", &full_hash[..8]);
        
        // 根據key類型返回不同長度：
        // AES-128-CTR key/IV: 16 bytes
        // HMAC-SHA256 key: 32 bytes
        match X {
            'A' | 'B' | 'C' | 'D' => Ok(full_hash[..16].to_vec()),  // IV or encryption key
            'E' | 'F' => Ok(full_hash.to_vec()),                   // MAC key (full 32 bytes)
            _ => Ok(full_hash[..16].to_vec()),                     // default
        }
    }
    
    /// SSH mpint编码（参考RFC 4253 Section 5）
    /// Curve25519 shared secret特殊处理
    fn encode_mpint(bytes: &[u8]) -> Vec<u8> {
        // RFC 4253: mpint = uint32(length) + data
        // 去掉前导零，如果最高位>=0x80前面加0
        
        // 去掉前导零字节（但不去掉最后一个字节即使它是0）
        let mut start = 0;
        while start < bytes.len() - 1 && bytes[start] == 0 {
            start += 1;
        }
        
        let data_without_leading_zeros = &bytes[start..];
        
        // 构建mpint数据
        let mut mpint_data = Vec::new();
        
        // 如果最高位>=0x80，前面加0字节（避免负数）
        if data_without_leading_zeros[0] >= 0x80 {
            mpint_data.push(0);
        }
        mpint_data.extend_from_slice(data_without_leading_zeros);
        
        // 最终格式：uint32长度 + mpint数据
        let mut result = Vec::new();
        result.extend_from_slice(&(mpint_data.len() as u32).to_be_bytes());
        result.extend_from_slice(&mpint_data);
        
        result
    }
}

/// Ed25519服务器主机密钥（参考OpenSSH sshkey.c）
pub struct Ed25519HostKey {
    signing_key: SigningKey,
}

impl Ed25519HostKey {
    /// 加载或生成主机密钥（参考OpenSSH hostfile.c）
    pub fn load_or_generate(key_path: &str) -> Result<Self> {
        // 简化实现：生成临时密钥（实际应从文件加载）
        // 参考OpenSSH ssh-keygen
        
        let signing_key = SigningKey::generate(&mut OsRng);
        
        Ok(Self { signing_key })
    }
    
    /// 获取公钥（用于SSH_MSG_KEX_ECDH_REPLY）
    pub fn public_key_bytes(&self) -> Vec<u8> {
        // SSH Ed25519公钥格式（参考OpenSSH sshkey.c）
        let verifying_key = self.signing_key.verifying_key();
        
        // SSH格式：ssh-ed25519 + 公钥bytes
        // 简化：仅返回公钥bytes（32字节）
        verifying_key.as_bytes().to_vec()
    }
    
    /// 签名（参考OpenSSH sshkey.c: sshkey_sign()）
    pub fn sign(&self, data: &[u8]) -> Result<Vec<u8>> {
        // OpenSSH Ed25519签名
        let signature = self.signing_key.sign(data);
        
        // SSH签名格式（参考OpenSSH ssh-sign.c）
        // 简化：仅返回签名bytes（64字节）
        Ok(signature.to_bytes().to_vec())
    }
    
    /// 获取完整SSH公钥格式（参考OpenSSH sshkey.c）
    pub fn ssh_public_key(&self) -> String {
        let public_bytes = self.public_key_bytes();
        
        // SSH公钥格式：ssh-ed25519 <base64-encoded-public-key>
        // 参考OpenSSH ssh-keygen -y
        
        use base64::{Engine as _, engine::general_purpose};
        let encoded = general_purpose::STANDARD.encode(&public_bytes);
        
        format!("ssh-ed25519 {}", encoded)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_curve25519_key_generation() {
        let kex = Curve25519Kex::new();
        assert_eq!(kex.public_key().len(), 32);
    }
    
    #[test]
    fn test_curve25519_shared_secret() {
        let client_kex = Curve25519Kex::new();
        let server_kex = Curve25519Kex::new();
        
        // 客户端计算共享密钥
        let client_secret = client_kex.compute_shared_secret(server_kex.public_key()).unwrap();
        
        // 服务器计算共享密钥
        let server_secret = server_kex.compute_shared_secret(client_kex.public_key()).unwrap();
        
        // 应该相同（Curve25519特性）
        assert_eq!(client_secret, server_secret);
    }
    
    #[test]
    fn test_ed25519_host_key() {
        let host_key = Ed25519HostKey::load_or_generate("test_key").unwrap();
        assert_eq!(host_key.public_key_bytes().len(), 32);
    }
    
    #[test]
    fn test_ed25519_signature() {
        let host_key = Ed25519HostKey::load_or_generate("test_key").unwrap();
        let data = b"test data";
        
        let signature = host_key.sign(data).unwrap();
        assert_eq!(signature.len(), 64);  // Ed25519签名64字节
    }
}