//! 分叉处理
//! 
//! 实现分叉检测、分叉选择、分叉恢复和分叉防范

use crate::block::Block;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

/// 分叉错误类型
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ForkError {
    /// 分叉检测失败
    DetectionFailed(String),
    /// 分叉选择失败
    SelectionFailed(String),
    /// 分叉恢复失败
    RecoveryFailed(String),
    /// 无效的分叉
    InvalidFork(String),
    /// 分叉链不存在
    ChainNotFound(String),
}

/// 分叉类型
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ForkType {
    /// 短期分叉（1-3个区块）
    ShortRange,
    /// 中期分叉（4-10个区块）
    MediumRange,
    /// 长期分叉（10+个区块）
    LongRange,
    /// 恶意分叉
    Malicious,
}

/// 分叉信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForkInfo {
    /// 分叉ID
    pub id: String,
    /// 分叉类型
    pub fork_type: ForkType,
    /// 分叉点高度
    pub fork_height: u64,
    /// 分叉链
    pub chains: Vec<ForkChain>,
    /// 检测时间
    pub detected_at: u64,
    /// 是否已解决
    pub resolved: bool,
}

impl ForkInfo {
    pub fn new(id: String, fork_height: u64) -> Self {
        ForkInfo {
            id,
            fork_type: ForkType::ShortRange,
            fork_height,
            chains: Vec::new(),
            detected_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .expect("FIX-006: unexpected None/Err")
                .as_secs(),
            resolved: false,
        }
    }

    /// 添加分叉链
    pub fn add_chain(&mut self, chain: ForkChain) {
        self.chains.push(chain);
        self.update_fork_type();
    }

    /// 更新分叉类型
    fn update_fork_type(&mut self) {
        let max_length = self.chains.iter().map(|c| c.length()).max().unwrap_or(0);
        
        self.fork_type = if max_length <= 3 {
            ForkType::ShortRange
        } else if max_length <= 10 {
            ForkType::MediumRange
        } else {
            ForkType::LongRange
        };
    }

    /// 标记为已解决
    pub fn mark_resolved(&mut self) {
        self.resolved = true;
    }
}

/// 分叉链
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForkChain {
    /// 链ID
    pub id: String,
    /// 区块列表
    pub blocks: Vec<Block>,
    /// 总权重
    pub total_weight: u64,
    /// 验证者集合
    pub producers: HashSet<String>,
}

impl ForkChain {
    pub fn new(id: String) -> Self {
        ForkChain {
            id,
            blocks: Vec::new(),
            total_weight: 0,
            producers: HashSet::new(),
        }
    }

    /// 添加区块
    pub fn add_block(&mut self, block: Block) {
        self.total_weight += 1; // 简化：每个区块权重为1
        self.producers.insert(block.header.producer.clone());
        self.blocks.push(block);
    }

    /// 获取链长度
    pub fn length(&self) -> usize {
        self.blocks.len()
    }

    /// 获取最新区块
    pub fn latest_block(&self) -> Option<&Block> {
        self.blocks.last()
    }

    /// 获取最新高度
    pub fn latest_height(&self) -> u64 {
        self.latest_block()
            .map(|b| b.header.height)
            .unwrap_or(0)
    }
}

/// 分叉检测器
#[derive(Debug)]
pub struct ForkDetector {
    /// 已知的分叉
    known_forks: HashMap<String, ForkInfo>,
    /// 区块索引（高度 -> 区块哈希列表）
    block_index: HashMap<u64, Vec<String>>,
    /// 区块存储
    block_store: HashMap<String, Block>,
    /// 检测阈值
    detection_threshold: usize,
}

impl ForkDetector {
    pub fn new(detection_threshold: usize) -> Self {
        ForkDetector {
            known_forks: HashMap::new(),
            block_index: HashMap::new(),
            block_store: HashMap::new(),
            detection_threshold,
        }
    }

    /// 添加区块
    pub fn add_block(&mut self, block: Block) -> Result<Option<ForkInfo>, ForkError> {
        let height = block.header.height;
        let hash = block.hash();
        
        // 存储区块
        self.block_store.insert(hash.clone(), block);
        
        // 更新索引
        let hashes = self.block_index.entry(height).or_insert_with(Vec::new);
        hashes.push(hash);
        
        // 检测分叉
        if hashes.len() > self.detection_threshold {
            let fork_id = format!("fork_{}_{}", height, hashes.len());
            let mut fork_info = ForkInfo::new(fork_id.clone(), height);
            
            // 构建分叉链
            for (i, h) in hashes.iter().enumerate() {
                if let Some(block) = self.block_store.get(h) {
                    let mut chain = ForkChain::new(format!("chain_{}_{}", height, i));
                    chain.add_block(block.clone());
                    fork_info.add_chain(chain);
                }
            }
            
            self.known_forks.insert(fork_id.clone(), fork_info.clone());
            return Ok(Some(fork_info));
        }
        
        Ok(None)
    }

    /// 获取分叉信息
    pub fn get_fork(&self, fork_id: &str) -> Option<&ForkInfo> {
        self.known_forks.get(fork_id)
    }

    /// 获取所有未解决的分叉
    pub fn get_unresolved_forks(&self) -> Vec<&ForkInfo> {
        self.known_forks
            .values()
            .filter(|f| !f.resolved)
            .collect()
    }

    /// 标记分叉已解决
    pub fn mark_fork_resolved(&mut self, fork_id: &str) -> Result<(), ForkError> {
        self.known_forks
            .get_mut(fork_id)
            .ok_or_else(|| ForkError::ChainNotFound(format!("Fork {} not found", fork_id)))?
            .mark_resolved();
        Ok(())
    }

    /// 获取统计信息
    pub fn stats(&self) -> ForkStats {
        let total_forks = self.known_forks.len();
        let resolved_forks = self.known_forks.values().filter(|f| f.resolved).count();
        let unresolved_forks = total_forks - resolved_forks;
        
        let mut fork_types = HashMap::new();
        for fork in self.known_forks.values() {
            *fork_types.entry(fork.fork_type.clone()).or_insert(0) += 1;
        }
        
        ForkStats {
            total_forks,
            resolved_forks,
            unresolved_forks,
            fork_types,
        }
    }
}

/// 分叉统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ForkStats {
    pub total_forks: usize,
    pub resolved_forks: usize,
    pub unresolved_forks: usize,
    pub fork_types: HashMap<ForkType, usize>,
}

/// 分叉选择规则
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ForkChoiceRule {
    /// 最长链规则
    LongestChain,
    /// 最重链规则（权重最大）
    HeaviestChain,
    /// GHOST规则（Greedy Heaviest Observed SubTree）
    Ghost,
    /// 最新区块规则
    LatestBlock,
}

/// 分叉选择器
#[derive(Debug)]
pub struct ForkChoiceSelector {
    /// 选择规则
    rule: ForkChoiceRule,
}

impl ForkChoiceSelector {
    pub fn new(rule: ForkChoiceRule) -> Self {
        ForkChoiceSelector { rule }
    }

    /// 选择最佳链
    pub fn select_best_chain<'a>(&self, fork_info: &'a ForkInfo) -> Result<&'a ForkChain, ForkError> {
        if fork_info.chains.is_empty() {
            return Err(ForkError::SelectionFailed(
                "No chains available for selection".to_string()
            ));
        }

        match self.rule {
            ForkChoiceRule::LongestChain => {
                fork_info.chains
                    .iter()
                    .max_by_key(|c| c.length())
                    .ok_or_else(|| ForkError::SelectionFailed("Failed to find longest chain".to_string()))
            }
            ForkChoiceRule::HeaviestChain => {
                fork_info.chains
                    .iter()
                    .max_by_key(|c| c.total_weight)
                    .ok_or_else(|| ForkError::SelectionFailed("Failed to find heaviest chain".to_string()))
            }
            ForkChoiceRule::Ghost => {
                // 简化实现：使用最重链
                fork_info.chains
                    .iter()
                    .max_by_key(|c| c.total_weight)
                    .ok_or_else(|| ForkError::SelectionFailed("Failed to apply GHOST rule".to_string()))
            }
            ForkChoiceRule::LatestBlock => {
                fork_info.chains
                    .iter()
                    .max_by_key(|c| c.latest_height())
                    .ok_or_else(|| ForkError::SelectionFailed("Failed to find latest block".to_string()))
            }
        }
    }

    /// 更新规则
    pub fn update_rule(&mut self, rule: ForkChoiceRule) {
        self.rule = rule;
    }

    /// 获取当前规则
    pub fn current_rule(&self) -> &ForkChoiceRule {
        &self.rule
    }
}

/// 分叉恢复器
#[derive(Debug)]
pub struct ForkRecovery {
    /// 恢复策略
    strategy: RecoveryStrategy,
    /// 最大回滚深度
    max_rollback_depth: u64,
}

impl ForkRecovery {
    pub fn new(strategy: RecoveryStrategy, max_rollback_depth: u64) -> Self {
        ForkRecovery {
            strategy,
            max_rollback_depth,
        }
    }

    /// 恢复分叉
    pub fn recover_from_fork(
        &self,
        fork_info: &ForkInfo,
        selected_chain: &ForkChain,
    ) -> Result<RecoveryPlan, ForkError> {
        match self.strategy {
            RecoveryStrategy::Rollback => {
                // 回滚到分叉点
                let rollback_depth = selected_chain.latest_height() - fork_info.fork_height;
                
                if rollback_depth > self.max_rollback_depth {
                    return Err(ForkError::RecoveryFailed(
                        format!("Rollback depth {} exceeds maximum {}", rollback_depth, self.max_rollback_depth)
                    ));
                }
                
                Ok(RecoveryPlan {
                    action: RecoveryAction::Rollback,
                    target_height: fork_info.fork_height,
                    blocks_to_apply: selected_chain.blocks.clone(),
                })
            }
            RecoveryStrategy::FastForward => {
                // 快进到最新区块
                Ok(RecoveryPlan {
                    action: RecoveryAction::FastForward,
                    target_height: selected_chain.latest_height(),
                    blocks_to_apply: selected_chain.blocks.clone(),
                })
            }
            RecoveryStrategy::Reorg => {
                // 重组区块链
                Ok(RecoveryPlan {
                    action: RecoveryAction::Reorg,
                    target_height: fork_info.fork_height,
                    blocks_to_apply: selected_chain.blocks.clone(),
                })
            }
        }
    }

    /// 更新策略
    pub fn update_strategy(&mut self, strategy: RecoveryStrategy) {
        self.strategy = strategy;
    }
}

/// 恢复策略
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum RecoveryStrategy {
    /// 回滚
    Rollback,
    /// 快进
    FastForward,
    /// 重组
    Reorg,
}

/// 恢复计划
#[derive(Debug, Clone)]
pub struct RecoveryPlan {
    /// 恢复动作
    pub action: RecoveryAction,
    /// 目标高度
    pub target_height: u64,
    /// 需要应用的区块
    pub blocks_to_apply: Vec<Block>,
}

/// 恢复动作
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RecoveryAction {
    /// 回滚
    Rollback,
    /// 快进
    FastForward,
    /// 重组
    Reorg,
}

/// 分叉防范器
#[derive(Debug)]
pub struct ForkPrevention {
    /// 最小区块生产者数量（OPN 开放生产网络要求）
    pub min_producers: usize,
    /// 最小 CR 权重阈值（用于分叉防范，不是投票权重）
    pub min_receipt_weight: u64,
    /// 黑名单验证者
    blacklisted_producers: HashSet<String>,
    /// 防范规则
    rules: Vec<PreventionRule>,
}

impl ForkPrevention {
    pub fn new(min_producers: usize, min_receipt_weight: u64) -> Self {
        ForkPrevention {
            min_producers,
            min_receipt_weight,
            blacklisted_producers: HashSet::new(),
            rules: Self::default_rules(),
        }
    }

    /// 默认防范规则
    fn default_rules() -> Vec<PreventionRule> {
        vec![
            PreventionRule {
                id: "rule_001".to_string(),
                name: "Minimum Block Producers".to_string(),
                description: "Require minimum number of CBP (Constitutional Block Producers) in OPN".to_string(),
                enabled: true,
            },
            PreventionRule {
                id: "rule_002".to_string(),
                name: "CR Receipt Weight Threshold".to_string(),
                description: "Require minimum cumulative CR receipt weight for fork prevention".to_string(),
                enabled: true,
            },
            PreventionRule {
                id: "rule_003".to_string(),
                name: "Blacklist Check".to_string(),
                description: "Block blacklisted producers (revoked DID/KYC)".to_string(),
                enabled: true,
            },
        ]
    }

    /// 检查区块是否可能导致分叉
    pub fn check_block(&self, block: &Block) -> Result<(), ForkError> {
        // 检查提议者是否在黑名单中
        if self.blacklisted_producers.contains(&block.header.producer) {
            return Err(ForkError::InvalidFork(
                format!("Proposer {} is blacklisted", block.header.producer)
            ));
        }

        // 检查区块签名数量
        // 简化实现：假设每个区块都有足够的签名
        
        Ok(())
    }

    /// 添加到黑名单
    pub fn add_to_blacklist(&mut self, producer: String) {
        self.blacklisted_producers.insert(producer);
    }

    /// 从黑名单移除
    pub fn remove_from_blacklist(&mut self, producer: &str) -> bool {
        self.blacklisted_producers.remove(producer)
    }

    /// 获取黑名单
    pub fn blacklist(&self) -> &HashSet<String> {
        &self.blacklisted_producers
    }

    /// 添加规则
    pub fn add_rule(&mut self, rule: PreventionRule) {
        self.rules.push(rule);
    }

    /// 获取所有规则
    pub fn rules(&self) -> &[PreventionRule] {
        &self.rules
    }
}

/// 防范规则
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PreventionRule {
    pub id: String,
    pub name: String,
    pub description: String,
    pub enabled: bool,
}

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

    #[test]
    fn test_fork_info_creation() {
        let fork = ForkInfo::new("test_fork".to_string(), 100);
        assert_eq!(fork.fork_height, 100);
        assert_eq!(fork.chains.len(), 0);
        assert!(!fork.resolved);
    }

    #[test]
    fn test_fork_chain() {
        let mut chain = ForkChain::new("chain1".to_string());
        let block = Block::new(1, "genesis".to_string(), "producer1".to_string());
        
        chain.add_block(block);
        assert_eq!(chain.length(), 1);
        assert_eq!(chain.total_weight, 1);
    }

    #[test]
    fn test_fork_detector() {
        let mut detector = ForkDetector::new(1);
        
        let block1 = Block::new(1, "genesis".to_string(), "producer1".to_string());
        let block2 = Block::new(1, "genesis".to_string(), "producer2".to_string());
        
        // 第一个区块不应该触发分叉
        assert!(detector.add_block(block1).expect("mainnet: handle error").is_none());
        
        // 第二个相同高度的区块应该触发分叉
        let fork = detector.add_block(block2).expect("FIX-006: unexpected None/Err");
        assert!(fork.is_some());
    }

    #[test]
    fn test_fork_choice_longest_chain() {
        let selector = ForkChoiceSelector::new(ForkChoiceRule::LongestChain);
        
        let mut fork_info = ForkInfo::new("test".to_string(), 1);
        
        let mut chain1 = ForkChain::new("chain1".to_string());
        chain1.add_block(Block::new(1, "genesis".to_string(), "v1".to_string()));
        
        let mut chain2 = ForkChain::new("chain2".to_string());
        chain2.add_block(Block::new(1, "genesis".to_string(), "v2".to_string()));
        chain2.add_block(Block::new(2, "block1".to_string(), "v2".to_string()));
        
        fork_info.add_chain(chain1);
        fork_info.add_chain(chain2);
        
        let best = selector.select_best_chain(&fork_info).expect("FIX-006: unexpected None/Err");
        assert_eq!(best.id, "chain2");
    }

    #[test]
    fn test_fork_recovery() {
        let recovery = ForkRecovery::new(RecoveryStrategy::Rollback, 100);
        
        let mut fork_info = ForkInfo::new("test".to_string(), 10);
        let mut chain = ForkChain::new("chain1".to_string());
        chain.add_block(Block::new(11, "block10".to_string(), "v1".to_string()));
        
        fork_info.add_chain(chain.clone());
        
        let plan = recovery.recover_from_fork(&fork_info, &chain).expect("FIX-006: unexpected None/Err");
        assert_eq!(plan.action, RecoveryAction::Rollback);
        assert_eq!(plan.target_height, 10);
    }

    #[test]
    fn test_fork_prevention() {
        let mut prevention = ForkPrevention::new(3, 1000);
        
        prevention.add_to_blacklist("malicious_producer".to_string());
        
        let block = Block::new(1, "genesis".to_string(), "malicious_producer".to_string());
        assert!(prevention.check_block(&block).is_err());
    }

    #[test]
    fn test_fork_stats() {
        let mut detector = ForkDetector::new(1);
        
        let block1 = Block::new(1, "genesis".to_string(), "v1".to_string());
        let block2 = Block::new(1, "genesis".to_string(), "v2".to_string());
        
        detector.add_block(block1).expect("FIX-006: unexpected None/Err");
        detector.add_block(block2).expect("FIX-006: unexpected None/Err");
        
        let stats = detector.stats();
        assert_eq!(stats.total_forks, 1);
        assert_eq!(stats.unresolved_forks, 1);
    }

    #[test]
    fn test_fork_type_update() {
        let mut fork_info = ForkInfo::new("test".to_string(), 1);
        
        let mut chain = ForkChain::new("chain1".to_string());
        for i in 1..=5 {
            chain.add_block(Block::new(i, format!("block{}", i-1), "v1".to_string()));
        }
        
        fork_info.add_chain(chain);
        assert_eq!(fork_info.fork_type, ForkType::MediumRange);
    }
}