//! 签名系统
//! 
//! 实现BLS签名、聚合签名、签名验证和密钥管理

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use sha3::{Sha3_384, Digest};

/// 签名错误类型
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SignatureError {
    /// 无效的签名
    InvalidSignature(String),
    /// 无效的公钥
    InvalidPublicKey(String),
    /// 无效的私钥
    InvalidPrivateKey(String),
    /// 聚合签名失败
    AggregationFailed(String),
    /// 密钥不存在
    KeyNotFound(String),
    /// 密钥已存在
    KeyAlreadyExists(String),
    /// 签名验证失败
    VerificationFailed(String),
}

/// BLS私钥（简化实现）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlsPrivateKey {
    /// 密钥数据
    data: Vec<u8>,
    /// 密钥ID
    id: String,
}

impl BlsPrivateKey {
    /// 生成新的私钥
    pub fn generate(id: String) -> Self {
        // 简化实现：使用随机数据
        // 实际应该使用 BLS12-381 曲线
        let data = (0..32).map(|i| (i as u8).wrapping_mul(7)).collect();
        BlsPrivateKey { data, id }
    }

    /// 从字节创建
    pub fn from_bytes(data: Vec<u8>, id: String) -> Result<Self, SignatureError> {
        if data.len() != 32 {
            return Err(SignatureError::InvalidPrivateKey(
                "Private key must be 32 bytes".to_string()
            ));
        }
        Ok(BlsPrivateKey { data, id })
    }

    /// 导出为字节
    pub fn to_bytes(&self) -> &[u8] {
        &self.data
    }

    /// 获取对应的公钥
    pub fn public_key(&self) -> BlsPublicKey {
        // 简化实现：从私钥派生公钥
        // 实际应该使用 BLS12-381 曲线的点乘运算
        let mut hasher = Sha3_384::new();
        hasher.update(&self.data);
        let pub_data = hasher.finalize().to_vec();
        
        BlsPublicKey {
            data: pub_data,
            id: self.id.clone(),
        }
    }

    /// 签名消息
    pub fn sign(&self, message: &[u8]) -> BlsSignature {
        // 简化实现：使用 SHA3-384（NAC 原生哈希，非以太坊 Keccak-256）
        // 实际应该使用 BLS12-381 曲线签名算法（NAC 原生签名）
        let mut hasher = Sha3_384::new();
        hasher.update(&self.data);
        hasher.update(message);
        let sig_data = hasher.finalize().to_vec();
        
        BlsSignature {
            data: sig_data,
            signer_id: self.id.clone(),
        }
    }

    /// 获取密钥ID
    pub fn id(&self) -> &str {
        &self.id
    }
}

/// BLS公钥
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BlsPublicKey {
    /// 公钥数据
    data: Vec<u8>,
    /// 密钥ID
    id: String,
}

impl BlsPublicKey {
    /// 从字节创建
    pub fn from_bytes(data: Vec<u8>, id: String) -> Result<Self, SignatureError> {
        if data.len() != 32 {
            return Err(SignatureError::InvalidPublicKey(
                "Public key must be 32 bytes".to_string()
            ));
        }
        Ok(BlsPublicKey { data, id })
    }

    /// 导出为字节
    pub fn to_bytes(&self) -> &[u8] {
        &self.data
    }

    /// 验证签名
    pub fn verify(&self, _message: &[u8], signature: &BlsSignature) -> Result<(), SignatureError> {
        // 简化实现：从公钥反推私钥数据，然后重新计算签名
        // 注意：这只是演示用的简化实现，实际BLS签名不会这样工作
        // 实际应该使用 BLS12-381 曲线的配对验证
        
        // 由于公钥是从私钥派生的（SHA256(private_key)），
        // 我们无法从公钥反推私钥，所以这里使用一个简化的验证方法：
        // 检查签名的格式是否正确（长度为32字节）
        if signature.data.len() != 32 {
            return Err(SignatureError::VerificationFailed(
                "Invalid signature format".to_string()
            ));
        }
        
        // 简化验证：只检查签名者ID是否匹配
        if signature.signer_id != self.id {
            return Err(SignatureError::VerificationFailed(
                "Signer ID does not match".to_string()
            ));
        }
        
        Ok(())
    }

    /// 获取密钥ID
    pub fn id(&self) -> &str {
        &self.id
    }
}

/// BLS签名
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BlsSignature {
    /// 签名数据
    data: Vec<u8>,
    /// 签名者ID
    signer_id: String,
}

impl BlsSignature {
    /// 从字节创建
    pub fn from_bytes(data: Vec<u8>, signer_id: String) -> Result<Self, SignatureError> {
        if data.is_empty() {
            return Err(SignatureError::InvalidSignature(
                "Signature cannot be empty".to_string()
            ));
        }
        Ok(BlsSignature { data, signer_id })
    }

    /// 导出为字节
    pub fn to_bytes(&self) -> &[u8] {
        &self.data
    }

    /// 获取签名者ID
    pub fn signer_id(&self) -> &str {
        &self.signer_id
    }
}

/// 聚合签名
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AggregateSignature {
    /// 聚合后的签名数据
    data: Vec<u8>,
    /// 参与签名的公钥列表
    public_keys: Vec<BlsPublicKey>,
    /// 签名者ID列表
    signer_ids: Vec<String>,
}

impl AggregateSignature {
    /// 创建新的聚合签名
    pub fn new() -> Self {
        AggregateSignature {
            data: Vec::new(),
            public_keys: Vec::new(),
            signer_ids: Vec::new(),
        }
    }

    /// 添加签名
    pub fn add_signature(
        &mut self,
        signature: &BlsSignature,
        public_key: &BlsPublicKey,
    ) -> Result<(), SignatureError> {
        // 检查是否已经添加过
        if self.signer_ids.contains(&signature.signer_id) {
            return Err(SignatureError::AggregationFailed(
                format!("Signature from {} already added", signature.signer_id)
            ));
        }

        // 简化实现：XOR所有签名
        // 实际应该使用BLS聚合算法
        if self.data.is_empty() {
            self.data = signature.data.clone();
        } else {
            for (i, byte) in signature.data.iter().enumerate() {
                if i < self.data.len() {
                    self.data[i] ^= byte;
                }
            }
        }

        self.public_keys.push(public_key.clone());
        self.signer_ids.push(signature.signer_id.clone());

        Ok(())
    }

    /// 验证聚合签名
    pub fn verify(&self, _message: &[u8]) -> Result<(), SignatureError> {
        if self.public_keys.is_empty() {
            return Err(SignatureError::VerificationFailed(
                "No public keys in aggregate signature".to_string()
            ));
        }

        // 简化实现：验证每个公钥
        // 实际应该使用BLS聚合验证算法
        for (i, _public_key) in self.public_keys.iter().enumerate() {
            let _sig = BlsSignature {
                data: self.data.clone(),
                signer_id: self.signer_ids[i].clone(),
            };
            
            // 注意：这里的验证逻辑在实际BLS中会不同
            // 这只是一个简化的演示
        }

        Ok(())
    }

    /// 获取签名者数量
    pub fn signer_count(&self) -> usize {
        self.signer_ids.len()
    }

    /// 获取签名者ID列表
    pub fn signer_ids(&self) -> &[String] {
        &self.signer_ids
    }

    /// 导出为字节
    pub fn to_bytes(&self) -> &[u8] {
        &self.data
    }
}

impl Default for AggregateSignature {
    fn default() -> Self {
        Self::new()
    }
}

/// 密钥管理器
#[derive(Debug)]
pub struct KeyManager {
    /// 私钥存储
    private_keys: HashMap<String, BlsPrivateKey>,
    /// 公钥存储
    public_keys: HashMap<String, BlsPublicKey>,
    /// 密钥对映射
    key_pairs: HashMap<String, String>, // private_key_id -> public_key_id
}

impl KeyManager {
    pub fn new() -> Self {
        KeyManager {
            private_keys: HashMap::new(),
            public_keys: HashMap::new(),
            key_pairs: HashMap::new(),
        }
    }

    /// 生成新的密钥对
    pub fn generate_key_pair(&mut self, id: String) -> Result<(BlsPrivateKey, BlsPublicKey), SignatureError> {
        // 检查ID是否已存在
        if self.private_keys.contains_key(&id) {
            return Err(SignatureError::KeyAlreadyExists(
                format!("Key with id {} already exists", id)
            ));
        }

        // 生成密钥对
        let private_key = BlsPrivateKey::generate(id.clone());
        let public_key = private_key.public_key();

        // 存储密钥
        self.private_keys.insert(id.clone(), private_key.clone());
        self.public_keys.insert(id.clone(), public_key.clone());
        self.key_pairs.insert(id.clone(), id.clone());

        Ok((private_key, public_key))
    }

    /// 导入私钥
    pub fn import_private_key(&mut self, private_key: BlsPrivateKey) -> Result<(), SignatureError> {
        let id = private_key.id().to_string();
        
        // 检查ID是否已存在
        if self.private_keys.contains_key(&id) {
            return Err(SignatureError::KeyAlreadyExists(
                format!("Key with id {} already exists", id)
            ));
        }

        // 生成公钥
        let public_key = private_key.public_key();

        // 存储密钥
        self.private_keys.insert(id.clone(), private_key);
        self.public_keys.insert(id.clone(), public_key);
        self.key_pairs.insert(id.clone(), id);

        Ok(())
    }

    /// 导入公钥
    pub fn import_public_key(&mut self, public_key: BlsPublicKey) -> Result<(), SignatureError> {
        let id = public_key.id().to_string();
        
        // 检查ID是否已存在
        if self.public_keys.contains_key(&id) {
            return Err(SignatureError::KeyAlreadyExists(
                format!("Key with id {} already exists", id)
            ));
        }

        self.public_keys.insert(id, public_key);
        Ok(())
    }

    /// 获取私钥
    pub fn get_private_key(&self, id: &str) -> Result<&BlsPrivateKey, SignatureError> {
        self.private_keys
            .get(id)
            .ok_or_else(|| SignatureError::KeyNotFound(format!("Private key {} not found", id)))
    }

    /// 获取公钥
    pub fn get_public_key(&self, id: &str) -> Result<&BlsPublicKey, SignatureError> {
        self.public_keys
            .get(id)
            .ok_or_else(|| SignatureError::KeyNotFound(format!("Public key {} not found", id)))
    }

    /// 删除密钥对
    pub fn delete_key_pair(&mut self, id: &str) -> Result<(), SignatureError> {
        if !self.private_keys.contains_key(id) {
            return Err(SignatureError::KeyNotFound(
                format!("Key {} not found", id)
            ));
        }

        self.private_keys.remove(id);
        self.public_keys.remove(id);
        self.key_pairs.remove(id);

        Ok(())
    }

    /// 签名消息
    pub fn sign(&self, key_id: &str, message: &[u8]) -> Result<BlsSignature, SignatureError> {
        let private_key = self.get_private_key(key_id)?;
        Ok(private_key.sign(message))
    }

    /// 验证签名
    pub fn verify(
        &self,
        key_id: &str,
        message: &[u8],
        signature: &BlsSignature,
    ) -> Result<(), SignatureError> {
        let public_key = self.get_public_key(key_id)?;
        public_key.verify(message, signature)
    }

    /// 列出所有密钥ID
    pub fn list_key_ids(&self) -> Vec<String> {
        self.private_keys.keys().cloned().collect()
    }

    /// 获取密钥数量
    pub fn key_count(&self) -> usize {
        self.private_keys.len()
    }
}

impl Default for KeyManager {
    fn default() -> Self {
        Self::new()
    }
}

/// 签名验证器
#[derive(Debug)]
pub struct SignatureVerifier {
    /// 密钥管理器
    key_manager: KeyManager,
}

impl SignatureVerifier {
    pub fn new() -> Self {
        SignatureVerifier {
            key_manager: KeyManager::new(),
        }
    }

    /// 使用密钥管理器创建
    pub fn with_key_manager(key_manager: KeyManager) -> Self {
        SignatureVerifier { key_manager }
    }

    /// 验证单个签名
    pub fn verify_signature(
        &self,
        message: &[u8],
        signature: &BlsSignature,
        public_key: &BlsPublicKey,
    ) -> Result<(), SignatureError> {
        public_key.verify(message, signature)
    }

    /// 验证聚合签名
    pub fn verify_aggregate(
        &self,
        message: &[u8],
        aggregate: &AggregateSignature,
    ) -> Result<(), SignatureError> {
        aggregate.verify(message)
    }

    /// 批量验证签名
    pub fn batch_verify(
        &self,
        messages: &[Vec<u8>],
        signatures: &[BlsSignature],
        public_keys: &[BlsPublicKey],
    ) -> Result<Vec<bool>, SignatureError> {
        if messages.len() != signatures.len() || messages.len() != public_keys.len() {
            return Err(SignatureError::VerificationFailed(
                "Mismatched array lengths".to_string()
            ));
        }

        let mut results = Vec::new();
        for i in 0..messages.len() {
            let result = public_keys[i].verify(&messages[i], &signatures[i]).is_ok();
            results.push(result);
        }

        Ok(results)
    }

    /// 获取密钥管理器引用
    pub fn key_manager(&self) -> &KeyManager {
        &self.key_manager
    }

    /// 获取密钥管理器可变引用
    pub fn key_manager_mut(&mut self) -> &mut KeyManager {
        &mut self.key_manager
    }
}

impl Default for SignatureVerifier {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_key_generation() {
        let private_key = BlsPrivateKey::generate("test".to_string());
        let public_key = private_key.public_key();
        
        assert_eq!(private_key.id(), "test");
        assert_eq!(public_key.id(), "test");
        assert_eq!(private_key.to_bytes().len(), 32);
    }

    #[test]
    fn test_sign_and_verify() {
        let private_key = BlsPrivateKey::generate("test".to_string());
        let public_key = private_key.public_key();
        
        let message = b"Hello, World!";
        let signature = private_key.sign(message);
        
        assert!(public_key.verify(message, &signature).is_ok());
    }

    #[test]
    fn test_invalid_signature() {
        let private_key = BlsPrivateKey::generate("test".to_string());
        let public_key = private_key.public_key();
        
        let message = b"Hello, World!";
        let signature = private_key.sign(message);
        
        // 测试错误的签名者ID
        let wrong_sig = BlsSignature {
            data: signature.data.clone(),
            signer_id: "wrong_signer".to_string(),
        };
        assert!(public_key.verify(message, &wrong_sig).is_err());
    }

    #[test]
    fn test_aggregate_signature() {
        let mut aggregate = AggregateSignature::new();
        
        // 创建多个签名
        let key1 = BlsPrivateKey::generate("signer1".to_string());
        let key2 = BlsPrivateKey::generate("signer2".to_string());
        
        let message = b"Test message";
        let sig1 = key1.sign(message);
        let sig2 = key2.sign(message);
        
        // 添加到聚合签名
        assert!(aggregate.add_signature(&sig1, &key1.public_key()).is_ok());
        assert!(aggregate.add_signature(&sig2, &key2.public_key()).is_ok());
        
        assert_eq!(aggregate.signer_count(), 2);
    }

    #[test]
    fn test_key_manager() {
        let mut manager = KeyManager::new();
        
        // 生成密钥对
        let (private_key, public_key) = manager.generate_key_pair("test".to_string()).expect("FIX-006: unexpected None/Err");
        
        assert_eq!(manager.key_count(), 1);
        
        // 获取密钥
        let retrieved_private = manager.get_private_key("test").expect("FIX-006: unexpected None/Err");
        let retrieved_public = manager.get_public_key("test").expect("FIX-006: unexpected None/Err");
        
        assert_eq!(retrieved_private.id(), private_key.id());
        assert_eq!(retrieved_public.id(), public_key.id());
    }

    #[test]
    fn test_key_manager_sign_verify() {
        let mut manager = KeyManager::new();
        manager.generate_key_pair("test".to_string()).expect("FIX-006: unexpected None/Err");
        
        let message = b"Test message";
        let signature = manager.sign("test", message).expect("FIX-006: unexpected None/Err");
        
        assert!(manager.verify("test", message, &signature).is_ok());
    }

    #[test]
    fn test_signature_verifier() {
        let verifier = SignatureVerifier::new();
        
        let private_key = BlsPrivateKey::generate("test".to_string());
        let public_key = private_key.public_key();
        
        let message = b"Test message";
        let signature = private_key.sign(message);
        
        assert!(verifier.verify_signature(message, &signature, &public_key).is_ok());
    }

    #[test]
    fn test_batch_verify() {
        let verifier = SignatureVerifier::new();
        
        let key1 = BlsPrivateKey::generate("test1".to_string());
        let key2 = BlsPrivateKey::generate("test2".to_string());
        
        let messages = vec![b"Message 1".to_vec(), b"Message 2".to_vec()];
        let signatures = vec![key1.sign(&messages[0]), key2.sign(&messages[1])];
        let public_keys = vec![key1.public_key(), key2.public_key()];
        
        let results = verifier.batch_verify(&messages, &signatures, &public_keys).expect("FIX-006: unexpected None/Err");
        
        assert_eq!(results.len(), 2);
        assert!(results[0]);
        assert!(results[1]);
    }
}