//! 多签名钱包模块
//!
//! 实现多签名地址、签名收集、签名验证和交易广播

use crate::WalletError;
use sha2::{Sha256, Digest};
use std::collections::{HashMap, HashSet};

/// 多签名方案类型
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MultisigScheme {
    /// M-of-N多签名（需要M个签名，共N个参与者）
    MOfN,
    /// 加权多签名（每个签名者有不同权重）
    Weighted,
    /// 分层多签名（多层审批）
    Hierarchical,
}

/// 签名者信息
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Signer {
    /// 签名者地址
    pub address: String,
    /// 公钥
    pub public_key: Vec<u8>,
    /// 权重（用于加权多签）
    pub weight: u32,
    /// 层级（用于分层多签）
    pub level: u32,
}

impl Signer {
    /// 创建新签名者
    pub fn new(address: String, public_key: Vec<u8>) -> Self {
        Signer {
            address,
            public_key,
            weight: 1,
            level: 0,
        }
    }

    /// 设置权重
    pub fn with_weight(mut self, weight: u32) -> Self {
        self.weight = weight;
        self
    }

    /// 设置层级
    pub fn with_level(mut self, level: u32) -> Self {
        self.level = level;
        self
    }
}

/// 多签名配置
#[derive(Debug, Clone)]
pub struct MultisigConfig {
    /// 方案类型
    pub scheme: MultisigScheme,
    /// 签名者列表
    pub signers: Vec<Signer>,
    /// 所需签名数（M-of-N中的M）
    pub required_signatures: usize,
    /// 所需权重阈值（加权多签）
    pub weight_threshold: u32,
    /// 超时时间（秒）
    pub timeout: u64,
}

impl MultisigConfig {
    /// 创建M-of-N多签配置
    pub fn m_of_n(signers: Vec<Signer>, required: usize) -> Result<Self, WalletError> {
        if required == 0 || required > signers.len() {
            return Err(WalletError::KeyError(
                format!("Invalid required signatures: {} (total: {})", required, signers.len())
            ));
        }

        Ok(MultisigConfig {
            scheme: MultisigScheme::MOfN,
            signers,
            required_signatures: required,
            weight_threshold: 0,
            timeout: 3600, // 默认1小时
        })
    }

    /// 创建加权多签配置
    pub fn weighted(signers: Vec<Signer>, threshold: u32) -> Result<Self, WalletError> {
        let total_weight: u32 = signers.iter().map(|s| s.weight).sum();
        if threshold == 0 || threshold > total_weight {
            return Err(WalletError::KeyError(
                format!("Invalid weight threshold: {} (total: {})", threshold, total_weight)
            ));
        }

        Ok(MultisigConfig {
            scheme: MultisigScheme::Weighted,
            signers,
            required_signatures: 0,
            weight_threshold: threshold,
            timeout: 3600,
        })
    }

    /// 创建分层多签配置
    pub fn hierarchical(signers: Vec<Signer>, required_per_level: HashMap<u32, usize>) -> Result<Self, WalletError> {
        // 验证每层的签名者数量
        let mut level_counts: HashMap<u32, usize> = HashMap::new();
        for signer in &signers {
            *level_counts.entry(signer.level).or_insert(0) += 1;
        }

        for (level, required) in &required_per_level {
            let count = level_counts.get(level).copied().unwrap_or(0);
            if *required > count {
                return Err(WalletError::KeyError(
                    format!("Level {} requires {} signatures but only has {} signers", 
                           level, required, count)
                ));
            }
        }

        Ok(MultisigConfig {
            scheme: MultisigScheme::Hierarchical,
            signers,
            required_signatures: required_per_level.values().sum(),
            weight_threshold: 0,
            timeout: 3600,
        })
    }

    /// 设置超时时间
    pub fn with_timeout(mut self, timeout: u64) -> Self {
        self.timeout = timeout;
        self
    }

    /// 验证签名者是否在配置中
    pub fn has_signer(&self, address: &str) -> bool {
        self.signers.iter().any(|s| s.address == address)
    }

    /// 获取签名者
    pub fn get_signer(&self, address: &str) -> Option<&Signer> {
        self.signers.iter().find(|s| s.address == address)
    }
}

/// 多签名地址
#[derive(Debug, Clone)]
pub struct MultisigAddress {
    /// 地址
    pub address: String,
    /// 配置
    pub config: MultisigConfig,
    /// 创建时间
    pub created_at: u64,
}

impl MultisigAddress {
    /// 从配置创建多签地址
    pub fn from_config(config: MultisigConfig, timestamp: u64) -> Self {
        let address = Self::generate_address(&config);
        MultisigAddress {
            address,
            config,
            created_at: timestamp,
        }
    }

    /// 生成多签地址
    fn generate_address(config: &MultisigConfig) -> String {
        let mut hasher = Sha256::new();
        
        // 添加方案类型
        hasher.update(&[config.scheme as u8]);
        
        // 添加所有签名者的公钥
        for signer in &config.signers {
            hasher.update(&signer.public_key);
        }
        
        // 添加阈值
        hasher.update(&config.required_signatures.to_be_bytes());
        hasher.update(&config.weight_threshold.to_be_bytes());
        
        let hash = hasher.finalize();
        format!("0x{}", hex::encode(&hash[..20]))
    }

    /// 获取地址
    pub fn address(&self) -> &str {
        &self.address
    }

    /// 获取配置
    pub fn config(&self) -> &MultisigConfig {
        &self.config
    }
}

/// 签名
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Signature {
    /// 签名者地址
    pub signer: String,
    /// 签名数据
    pub data: Vec<u8>,
    /// 签名时间
    pub timestamp: u64,
}

impl Signature {
    /// 创建新签名
    pub fn new(signer: String, data: Vec<u8>, timestamp: u64) -> Self {
        Signature {
            signer,
            data,
            timestamp,
        }
    }
}

/// 多签名交易
#[derive(Debug, Clone)]
pub struct MultisigTransaction {
    /// 交易ID
    pub id: String,
    /// 多签地址
    pub multisig_address: String,
    /// 交易数据
    pub transaction_data: Vec<u8>,
    /// 已收集的签名
    pub signatures: Vec<Signature>,
    /// 创建时间
    pub created_at: u64,
    /// 过期时间
    pub expires_at: u64,
    /// 状态
    pub status: TransactionStatus,
}

/// 交易状态
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransactionStatus {
    /// 待签名
    Pending,
    /// 已完成
    Completed,
    /// 已广播
    Broadcasted,
    /// 已过期
    Expired,
    /// 已拒绝
    Rejected,
}

impl MultisigTransaction {
    /// 创建新的多签交易
    pub fn new(
        multisig_address: String,
        transaction_data: Vec<u8>,
        timestamp: u64,
        timeout: u64
    ) -> Self {
        let id = Self::generate_id(&multisig_address, &transaction_data, timestamp);
        
        MultisigTransaction {
            id,
            multisig_address,
            transaction_data,
            signatures: Vec::new(),
            created_at: timestamp,
            expires_at: timestamp + timeout,
            status: TransactionStatus::Pending,
        }
    }

    /// 生成交易ID
    fn generate_id(address: &str, data: &[u8], timestamp: u64) -> String {
        let mut hasher = Sha256::new();
        hasher.update(address.as_bytes());
        hasher.update(data);
        hasher.update(&timestamp.to_be_bytes());
        let hash = hasher.finalize();
        hex::encode(hash)
    }

    /// 添加签名
    pub fn add_signature(&mut self, signature: Signature) -> Result<(), WalletError> {
        // 检查状态
        if self.status != TransactionStatus::Pending {
            return Err(WalletError::Other(
                format!("Transaction is not pending: {:?}", self.status)
            ));
        }

        // 检查是否已签名
        if self.signatures.iter().any(|s| s.signer == signature.signer) {
            return Err(WalletError::Other(
                format!("Signer {} has already signed", signature.signer)
            ));
        }

        self.signatures.push(signature);
        Ok(())
    }

    /// 检查是否过期
    pub fn is_expired(&self, current_time: u64) -> bool {
        current_time > self.expires_at
    }

    /// 获取签名者列表
    pub fn get_signers(&self) -> HashSet<String> {
        self.signatures.iter()
            .map(|s| s.signer.clone())
            .collect()
    }
}

/// 签名收集器
pub struct SignatureCollector {
    /// 待签名交易
    transactions: HashMap<String, MultisigTransaction>,
    /// 多签地址配置
    addresses: HashMap<String, MultisigAddress>,
}

impl SignatureCollector {
    /// 创建新的签名收集器
    pub fn new() -> Self {
        SignatureCollector {
            transactions: HashMap::new(),
            addresses: HashMap::new(),
        }
    }

    /// 注册多签地址
    pub fn register_address(&mut self, address: MultisigAddress) {
        self.addresses.insert(address.address.clone(), address);
    }

    /// 创建交易
    pub fn create_transaction(
        &mut self,
        multisig_address: String,
        transaction_data: Vec<u8>,
        timestamp: u64
    ) -> Result<String, WalletError> {
        // 验证多签地址
        let address_config = self.addresses.get(&multisig_address)
            .ok_or_else(|| WalletError::Other("Multisig address not found".to_string()))?;

        let timeout = address_config.config.timeout;
        let tx = MultisigTransaction::new(multisig_address, transaction_data, timestamp, timeout);
        let tx_id = tx.id.clone();
        
        self.transactions.insert(tx_id.clone(), tx);
        Ok(tx_id)
    }

    /// 添加签名
    pub fn add_signature(
        &mut self,
        tx_id: &str,
        signature: Signature,
        current_time: u64
    ) -> Result<(), WalletError> {
        let tx = self.transactions.get_mut(tx_id)
            .ok_or_else(|| WalletError::Other("Transaction not found".to_string()))?;

        // 检查过期
        if tx.is_expired(current_time) {
            tx.status = TransactionStatus::Expired;
            return Err(WalletError::Other("Transaction expired".to_string()));
        }

        // 验证签名者
        let address_config = self.addresses.get(&tx.multisig_address)
            .ok_or_else(|| WalletError::Other("Multisig address not found".to_string()))?;

        if !address_config.config.has_signer(&signature.signer) {
            return Err(WalletError::Other(
                format!("Signer {} is not authorized", signature.signer)
            ));
        }

        tx.add_signature(signature)?;
        Ok(())
    }

    /// 检查交易是否可以广播
    pub fn can_broadcast(&self, tx_id: &str) -> Result<bool, WalletError> {
        let tx = self.transactions.get(tx_id)
            .ok_or_else(|| WalletError::Other("Transaction not found".to_string()))?;

        let address_config = self.addresses.get(&tx.multisig_address)
            .ok_or_else(|| WalletError::Other("Multisig address not found".to_string()))?;

        Ok(self.verify_signatures(tx, &address_config.config))
    }

    /// 验证签名
    fn verify_signatures(&self, tx: &MultisigTransaction, config: &MultisigConfig) -> bool {
        match config.scheme {
            MultisigScheme::MOfN => {
                tx.signatures.len() >= config.required_signatures
            }
            MultisigScheme::Weighted => {
                let total_weight: u32 = tx.signatures.iter()
                    .filter_map(|sig| config.get_signer(&sig.signer))
                    .map(|signer| signer.weight)
                    .sum();
                total_weight >= config.weight_threshold
            }
            MultisigScheme::Hierarchical => {
                // 检查每层是否满足要求
                let mut level_counts: HashMap<u32, usize> = HashMap::new();
                for sig in &tx.signatures {
                    if let Some(signer) = config.get_signer(&sig.signer) {
                        *level_counts.entry(signer.level).or_insert(0) += 1;
                    }
                }
                
                // 简化实现：假设每层需要至少1个签名
                let max_level = config.signers.iter().map(|s| s.level).max().unwrap_or(0);
                for level in 0..=max_level {
                    if level_counts.get(&level).copied().unwrap_or(0) == 0 {
                        return false;
                    }
                }
                true
            }
        }
    }

    /// 广播交易
    pub fn broadcast_transaction(&mut self, tx_id: &str) -> Result<(), WalletError> {
        if !self.can_broadcast(tx_id)? {
            return Err(WalletError::Other("Insufficient signatures".to_string()));
        }

        let tx = self.transactions.get_mut(tx_id)
            .ok_or_else(|| WalletError::Other("Transaction not found".to_string()))?;

        tx.status = TransactionStatus::Broadcasted;
        Ok(())
    }

    /// 获取交易
    pub fn get_transaction(&self, tx_id: &str) -> Option<&MultisigTransaction> {
        self.transactions.get(tx_id)
    }

    /// 获取待签名交易列表
    pub fn get_pending_transactions(&self, signer_address: &str) -> Vec<&MultisigTransaction> {
        self.transactions.values()
            .filter(|tx| {
                tx.status == TransactionStatus::Pending &&
                !tx.get_signers().contains(signer_address)
            })
            .collect()
    }

    /// 清理过期交易
    pub fn cleanup_expired(&mut self, current_time: u64) {
        let expired_ids: Vec<String> = self.transactions.iter()
            .filter(|(_, tx)| tx.is_expired(current_time) && tx.status == TransactionStatus::Pending)
            .map(|(id, _)| id.clone())
            .collect();

        for id in expired_ids {
            if let Some(tx) = self.transactions.get_mut(&id) {
                tx.status = TransactionStatus::Expired;
            }
        }
    }
}

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

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

    fn create_test_signer(index: u8) -> Signer {
        Signer::new(
            format!("0x{:040x}", index),
            vec![index; 32]
        )
    }

    #[test]
    fn test_multisig_config_m_of_n() {
        let signers = vec![
            create_test_signer(1),
            create_test_signer(2),
            create_test_signer(3),
        ];
        let config = MultisigConfig::m_of_n(signers, 2).unwrap();
        assert_eq!(config.scheme, MultisigScheme::MOfN);
        assert_eq!(config.required_signatures, 2);
    }

    #[test]
    fn test_multisig_config_weighted() {
        let signers = vec![
            create_test_signer(1).with_weight(2),
            create_test_signer(2).with_weight(3),
            create_test_signer(3).with_weight(5),
        ];
        let config = MultisigConfig::weighted(signers, 7).unwrap();
        assert_eq!(config.scheme, MultisigScheme::Weighted);
        assert_eq!(config.weight_threshold, 7);
    }

    #[test]
    fn test_multisig_address_generation() {
        let signers = vec![
            create_test_signer(1),
            create_test_signer(2),
        ];
        let config = MultisigConfig::m_of_n(signers, 2).unwrap();
        let address = MultisigAddress::from_config(config, 1000);
        assert!(address.address.starts_with("0x"));
        assert_eq!(address.address.len(), 42);
    }

    #[test]
    fn test_multisig_transaction_creation() {
        let tx = MultisigTransaction::new(
            "0x1234".to_string(),
            vec![1, 2, 3],
            1000,
            3600
        );
        assert_eq!(tx.status, TransactionStatus::Pending);
        assert_eq!(tx.signatures.len(), 0);
        assert_eq!(tx.expires_at, 4600);
    }

    #[test]
    fn test_add_signature() {
        let mut tx = MultisigTransaction::new(
            "0x1234".to_string(),
            vec![1, 2, 3],
            1000,
            3600
        );
        let sig = Signature::new("0x5678".to_string(), vec![4, 5, 6], 1100);
        tx.add_signature(sig).unwrap();
        assert_eq!(tx.signatures.len(), 1);
    }

    #[test]
    fn test_duplicate_signature() {
        let mut tx = MultisigTransaction::new(
            "0x1234".to_string(),
            vec![1, 2, 3],
            1000,
            3600
        );
        let sig1 = Signature::new("0x5678".to_string(), vec![4, 5, 6], 1100);
        let sig2 = Signature::new("0x5678".to_string(), vec![7, 8, 9], 1200);
        tx.add_signature(sig1).unwrap();
        let result = tx.add_signature(sig2);
        assert!(result.is_err());
    }

    #[test]
    fn test_signature_collector() {
        let mut collector = SignatureCollector::new();
        
        let signers = vec![
            create_test_signer(1),
            create_test_signer(2),
        ];
        let config = MultisigConfig::m_of_n(signers, 2).unwrap();
        let address = MultisigAddress::from_config(config, 1000);
        collector.register_address(address.clone());

        let tx_id = collector.create_transaction(
            address.address.clone(),
            vec![1, 2, 3],
            1000
        ).unwrap();

        assert!(collector.get_transaction(&tx_id).is_some());
    }

    #[test]
    fn test_can_broadcast_m_of_n() {
        let mut collector = SignatureCollector::new();
        
        let signers = vec![
            create_test_signer(1),
            create_test_signer(2),
            create_test_signer(3),
        ];
        let config = MultisigConfig::m_of_n(signers.clone(), 2).unwrap();
        let address = MultisigAddress::from_config(config, 1000);
        collector.register_address(address.clone());

        let tx_id = collector.create_transaction(
            address.address.clone(),
            vec![1, 2, 3],
            1000
        ).unwrap();

        // 添加第一个签名
        let sig1 = Signature::new(signers[0].address.clone(), vec![1], 1100);
        collector.add_signature(&tx_id, sig1, 1100).unwrap();
        assert!(!collector.can_broadcast(&tx_id).unwrap());

        // 添加第二个签名
        let sig2 = Signature::new(signers[1].address.clone(), vec![2], 1200);
        collector.add_signature(&tx_id, sig2, 1200).unwrap();
        assert!(collector.can_broadcast(&tx_id).unwrap());
    }

    #[test]
    fn test_broadcast_transaction() {
        let mut collector = SignatureCollector::new();
        
        let signers = vec![
            create_test_signer(1),
            create_test_signer(2),
        ];
        let config = MultisigConfig::m_of_n(signers.clone(), 2).unwrap();
        let address = MultisigAddress::from_config(config, 1000);
        collector.register_address(address.clone());

        let tx_id = collector.create_transaction(
            address.address.clone(),
            vec![1, 2, 3],
            1000
        ).unwrap();

        let sig1 = Signature::new(signers[0].address.clone(), vec![1], 1100);
        let sig2 = Signature::new(signers[1].address.clone(), vec![2], 1200);
        collector.add_signature(&tx_id, sig1, 1100).unwrap();
        collector.add_signature(&tx_id, sig2, 1200).unwrap();

        collector.broadcast_transaction(&tx_id).unwrap();
        let tx = collector.get_transaction(&tx_id).unwrap();
        assert_eq!(tx.status, TransactionStatus::Broadcasted);
    }

    #[test]
    fn test_expired_transaction() {
        let mut tx = MultisigTransaction::new(
            "0x1234".to_string(),
            vec![1, 2, 3],
            1000,
            100
        );
        assert!(!tx.is_expired(1050));
        assert!(tx.is_expired(1150));
    }

    #[test]
    fn test_cleanup_expired() {
        let mut collector = SignatureCollector::new();
        
        let signers = vec![create_test_signer(1)];
        let config = MultisigConfig::m_of_n(signers, 1).unwrap();
        let address = MultisigAddress::from_config(config, 1000);
        collector.register_address(address.clone());

        let tx_id = collector.create_transaction(
            address.address.clone(),
            vec![1, 2, 3],
            1000
        ).unwrap();

        collector.cleanup_expired(5000);
        let tx = collector.get_transaction(&tx_id).unwrap();
        assert_eq!(tx.status, TransactionStatus::Expired);
    }
}