// nac-cbpp-l0/src/jurisdiction_block.rs
// CBPP 多辖区区块头扩展
//
// CBPP 四大原则：
//   约法即是治法 — 宪法条款直接约束区块生产行为，无宪法外治理
//   宪法即是规则 — 区块合法性由宪法条款判定，规则来源于辖区真实法律
//   参与即是共识 — 节点将交易打包进区块，即是对所有涉及辖区 CR 的背书，无需多签
//   节点产生区块，交易决定区块大小 — jurisdiction_merkle_root 大小由实际
//     跨辖区交易数量动态决定，不是固定字段

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

/// 辖区 ID（ISO 3166-1 数字码 + 子辖区扩展）
/// 高 16 位：国家码；低 16 位：子辖区扩展（0 = 全国）
pub type JurisdictionId = u32;

/// 辖区宪法收据哈希（SHA3-384，48 字节）
///
/// 由 CEE 节点执行辖区插件后独立生成。
/// 辖区插件的规则来源于辖区政府授权 CA 签名的真实法律文件，
/// 不是链上治理决定。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CrHash(pub Vec<u8>);

impl CrHash {
    pub fn new(bytes: [u8; 48]) -> Self {
        Self(bytes.to_vec())
    }

    pub fn zero() -> Self {
        Self(vec![0u8; 48])
    }

    pub fn is_valid(&self) -> bool {
        self.0.len() == 48 && self.0 != vec![0u8; 48]
    }
}

/// 单笔跨辖区交易的辖区收据集合
///
/// 每个涉及辖区都有一个独立的 CR 哈希。
/// 各辖区 CEE 节点独立执行本辖区插件后生成，
/// 节点将其打包进区块即是背书（参与即共识）。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TxJurisdictionReceipts {
    /// 交易 ID（SHA3-384，48 字节）
    pub tx_id: Vec<u8>,
    /// 辖区 → CR 哈希 映射（BTreeMap 保证确定性排序）
    pub receipts: BTreeMap<JurisdictionId, CrHash>,
}

impl TxJurisdictionReceipts {
    pub fn new(tx_id: Vec<u8>) -> Self {
        Self {
            tx_id,
            receipts: BTreeMap::new(),
        }
    }

    /// 添加某辖区的独立 CR
    /// 各辖区 CEE 节点执行本辖区插件后直接写入，无需其他辖区确认
    pub fn add_receipt(&mut self, jurisdiction: JurisdictionId, cr_hash: CrHash) {
        self.receipts.insert(jurisdiction, cr_hash);
    }

    /// 验证指定辖区的 CR 是否存在且有效
    pub fn has_valid_receipt_for(&self, jurisdiction: JurisdictionId) -> bool {
        self.receipts
            .get(&jurisdiction)
            .map(|cr| cr.is_valid())
            .unwrap_or(false)
    }

    /// 获取所有涉及辖区列表
    pub fn jurisdictions(&self) -> Vec<JurisdictionId> {
        self.receipts.keys().copied().collect()
    }
}

/// 区块辖区默克尔树
///
/// 由区块内所有跨辖区交易的 CR 哈希构建。
/// 区块大小由实际交易数量决定，默克尔树大小随之动态变化。
/// 节点将此树根写入区块头，即是对所有辖区规则合规性的背书。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JurisdictionMerkleTree {
    /// 叶节点列表
    pub leaves: Vec<Vec<u8>>,
    /// 默克尔树根（48 字节 SHA3-384）
    pub root: Vec<u8>,
    /// 涉及辖区总数
    pub jurisdiction_count: u32,
    /// 跨辖区交易总数（体现"交易决定区块大小"原则）
    pub cross_tx_count: u32,
}

impl JurisdictionMerkleTree {
    /// 从区块内所有跨辖区交易的收据集合构建默克尔树
    /// 交易数量决定树的大小，完全动态
    pub fn build(tx_receipts: &[TxJurisdictionReceipts]) -> Self {
        let mut leaves: Vec<Vec<u8>> = Vec::new();
        let mut all_jurisdictions = BTreeSet::new();

        for tx_receipt in tx_receipts {
            for (jid, cr_hash) in &tx_receipt.receipts {
                // 叶节点 = SHA3-384(tx_id || jid_be_bytes || cr_hash)
                let mut hasher = Sha3_384::new();
                hasher.update(&tx_receipt.tx_id);
                hasher.update(&jid.to_be_bytes());
                hasher.update(&cr_hash.0);
                leaves.push(hasher.finalize().to_vec());
                all_jurisdictions.insert(*jid);
            }
        }

        let root = if leaves.is_empty() {
            vec![0u8; 48]
        } else {
            Self::compute_root(&leaves)
        };

        Self {
            jurisdiction_count: all_jurisdictions.len() as u32,
            cross_tx_count: tx_receipts.len() as u32,
            root,
            leaves,
        }
    }

    /// SHA3-384 两两合并默克尔树根
    fn compute_root(leaves: &[Vec<u8>]) -> Vec<u8> {
        if leaves.len() == 1 {
            return leaves[0].clone();
        }
        let mut level = leaves.to_vec();
        while level.len() > 1 {
            let mut next = Vec::new();
            let mut i = 0;
            while i < level.len() {
                let left = &level[i];
                let right = level.get(i + 1).unwrap_or(&level[i]);
                let mut h = Sha3_384::new();
                h.update(left);
                h.update(right);
                next.push(h.finalize().to_vec());
                i += 2;
            }
            level = next;
        }
        level.into_iter().next().unwrap_or_else(|| vec![0u8; 48])
    }

    /// 验证树根是否与期望值一致
    pub fn verify_root(&self, expected: &[u8]) -> bool {
        self.root == expected
    }
}

/// CBPP 多辖区区块头扩展字段
///
/// 嵌入标准区块头，记录该区块内所有跨辖区交易的合规证明。
/// 节点将此扩展写入区块头即完成对所有涉及辖区宪法规则的背书（参与即共识）。
/// 区块大小由实际交易数量决定，此扩展大小随之动态变化。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlockJurisdictionExtension {
    /// 辖区默克尔树根（48 字节 SHA3-384）
    /// 若区块内无跨辖区交易，则为全零
    pub jurisdiction_merkle_root: Vec<u8>,
    /// 涉及辖区数量
    pub jurisdiction_count: u32,
    /// 跨辖区交易数量（体现"交易决定区块大小"）
    pub cross_tx_count: u32,
    /// 宪法版本哈希（辖区政府授权 CA 签名的法律文件集合哈希）
    /// 确保所有节点使用相同版本的真实法律规则出块
    pub constitution_version_hash: Vec<u8>,
}

impl BlockJurisdictionExtension {
    /// 从跨辖区交易收据集合构建区块头扩展
    pub fn build(
        tx_receipts: &[TxJurisdictionReceipts],
        constitution_version_hash: Vec<u8>,
    ) -> Self {
        if tx_receipts.is_empty() {
            return Self {
                jurisdiction_merkle_root: vec![0u8; 48],
                jurisdiction_count: 0,
                cross_tx_count: 0,
                constitution_version_hash,
            };
        }

        let tree = JurisdictionMerkleTree::build(tx_receipts);
        Self {
            jurisdiction_merkle_root: tree.root,
            jurisdiction_count: tree.jurisdiction_count,
            cross_tx_count: tree.cross_tx_count,
            constitution_version_hash,
        }
    }

    /// 验证区块头扩展的合法性
    pub fn validate(&self, tx_receipts: &[TxJurisdictionReceipts]) -> BlockExtensionValidation {
        // 纯链内交易区块
        if self.cross_tx_count == 0 {
            return if self.jurisdiction_merkle_root == vec![0u8; 48] {
                BlockExtensionValidation::Valid
            } else {
                BlockExtensionValidation::Invalid(
                    "跨辖区交易数为0但默克尔根非零".to_string(),
                )
            };
        }

        // 重新计算默克尔树验证
        let tree = JurisdictionMerkleTree::build(tx_receipts);

        if tree.root != self.jurisdiction_merkle_root {
            return BlockExtensionValidation::Invalid(
                "jurisdiction_merkle_root 与交易收据不匹配".to_string(),
            );
        }
        if tree.jurisdiction_count != self.jurisdiction_count {
            return BlockExtensionValidation::Invalid(
                "辖区数量与默克尔树不一致".to_string(),
            );
        }
        if tree.cross_tx_count != self.cross_tx_count {
            return BlockExtensionValidation::Invalid(
                "跨辖区交易数量与默克尔树不一致".to_string(),
            );
        }

        BlockExtensionValidation::Valid
    }

    /// 是否为纯链内交易区块
    pub fn is_intra_chain_only(&self) -> bool {
        self.cross_tx_count == 0
    }
}

/// 区块头扩展验证结果
#[derive(Debug, PartialEq)]
pub enum BlockExtensionValidation {
    Valid,
    Invalid(String),
}

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

    fn tx_id(n: u8) -> Vec<u8> { vec![n; 48] }
    fn cr(n: u8) -> CrHash { CrHash::new([n; 48]) }

    #[test]
    fn test_single_cross_tx_two_jurisdictions() {
        let mut r = TxJurisdictionReceipts::new(tx_id(1));
        r.add_receipt(784, cr(10)); // UAE
        r.add_receipt(156, cr(20)); // China

        let ext = BlockJurisdictionExtension::build(&[r.clone()], vec![0xAB; 48]);
        assert_eq!(ext.cross_tx_count, 1);
        assert_eq!(ext.jurisdiction_count, 2);
        assert_ne!(ext.jurisdiction_merkle_root, vec![0u8; 48]);
        assert_eq!(ext.validate(&[r]), BlockExtensionValidation::Valid);
    }

    #[test]
    fn test_intra_chain_block() {
        let ext = BlockJurisdictionExtension::build(&[], vec![0xAB; 48]);
        assert!(ext.is_intra_chain_only());
        assert_eq!(ext.validate(&[]), BlockExtensionValidation::Valid);
    }

    #[test]
    fn test_dynamic_size_by_tx_count() {
        // 交易决定区块大小：10笔跨辖区交易 vs 1笔
        let mut receipts_10 = Vec::new();
        for i in 0..10u8 {
            let mut r = TxJurisdictionReceipts::new(tx_id(i));
            r.add_receipt(784, cr(i));
            r.add_receipt(156, cr(i + 50));
            receipts_10.push(r);
        }
        let ext_10 = BlockJurisdictionExtension::build(&receipts_10, vec![1u8; 48]);

        let mut r1 = TxJurisdictionReceipts::new(tx_id(0));
        r1.add_receipt(784, cr(0));
        r1.add_receipt(156, cr(50));
        let ext_1 = BlockJurisdictionExtension::build(&[r1], vec![1u8; 48]);

        // 10笔交易的默克尔根与1笔不同，体现动态性
        assert_ne!(ext_10.jurisdiction_merkle_root, ext_1.jurisdiction_merkle_root);
        assert_eq!(ext_10.cross_tx_count, 10);
        assert_eq!(ext_1.cross_tx_count, 1);

        assert_eq!(ext_10.validate(&receipts_10), BlockExtensionValidation::Valid);
    }

    #[test]
    fn test_tampered_receipt_rejected() {
        let mut r = TxJurisdictionReceipts::new(tx_id(1));
        r.add_receipt(784, cr(10));
        r.add_receipt(156, cr(20));
        let ext = BlockJurisdictionExtension::build(&[r], vec![0xAB; 48]);

        // 篡改 CR
        let mut tampered = TxJurisdictionReceipts::new(tx_id(1));
        tampered.add_receipt(784, cr(99));
        tampered.add_receipt(156, cr(20));
        assert!(matches!(ext.validate(&[tampered]), BlockExtensionValidation::Invalid(_)));
    }

    #[test]
    fn test_participation_is_consensus_no_multisig() {
        // 参与即共识：节点将 receipts 打包进区块头即完成背书，无需多签
        let mut r = TxJurisdictionReceipts::new(tx_id(42));
        r.add_receipt(784, cr(1)); // 单节点独立写入，无需其他节点确认
        let ext = BlockJurisdictionExtension::build(&[r.clone()], vec![1u8; 48]);
        // 验证通过即代表共识达成
        assert_eq!(ext.validate(&[r]), BlockExtensionValidation::Valid);
    }

    #[test]
    fn test_constitution_version_is_real_law_hash() {
        // 宪法版本哈希代表辖区政府授权CA签名的真实法律文件集合
        // 不是链上治理投票结果
        let law_hash = vec![0xDE; 48]; // 模拟真实法律文件哈希
        let ext = BlockJurisdictionExtension::build(&[], law_hash.clone());
        assert_eq!(ext.constitution_version_hash, law_hash);
    }
}