//! Merkle proof aggregation for efficient cross-shard verification.
//!
//! Aggregates state proofs from multiple shards into a single verifiable proof.

use serde::{Deserialize, Serialize};
use synor_types::Hash256;

use crate::{
    state::{ShardStateManager, StateProof},
    ShardError, ShardId, ShardResult,
};

/// Aggregated proof from multiple shards.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AggregatedProof {
    /// Beacon chain block this proof is for.
    pub beacon_block: Hash256,
    /// Per-shard state roots.
    pub shard_roots: Vec<(ShardId, Hash256)>,
    /// Combined root hash (Merkle root of shard roots).
    pub combined_root: Hash256,
    /// Individual proofs (optional, for specific account lookups).
    pub proofs: Vec<StateProof>,
    /// Timestamp when proof was generated.
    pub timestamp: u64,
}

impl AggregatedProof {
    /// Verifies the aggregated proof.
    pub fn verify(&self) -> bool {
        // Verify combined root matches shard roots
        let computed_root = Self::compute_combined_root(&self.shard_roots);
        computed_root == self.combined_root
    }

    /// Computes the combined root from shard roots.
    fn compute_combined_root(shard_roots: &[(ShardId, Hash256)]) -> Hash256 {
        if shard_roots.is_empty() {
            return Hash256::from_bytes([0u8; 32]);
        }

        // Build Merkle tree from shard roots
        let mut hashes: Vec<Hash256> = shard_roots
            .iter()
            .map(|(id, root)| {
                let mut hasher = blake3::Hasher::new();
                hasher.update(&id.to_le_bytes());
                hasher.update(root.as_bytes());
                Hash256::from_bytes(*hasher.finalize().as_bytes())
            })
            .collect();

        // Pad to power of 2
        let target_len = hashes.len().next_power_of_two();
        while hashes.len() < target_len {
            hashes.push(Hash256::from_bytes([0u8; 32]));
        }

        // Build tree bottom-up
        while hashes.len() > 1 {
            let mut next_level = Vec::new();
            for chunk in hashes.chunks(2) {
                let mut hasher = blake3::Hasher::new();
                hasher.update(chunk[0].as_bytes());
                if chunk.len() > 1 {
                    hasher.update(chunk[1].as_bytes());
                }
                next_level.push(Hash256::from_bytes(*hasher.finalize().as_bytes()));
            }
            hashes = next_level;
        }

        hashes[0]
    }

    /// Gets the state root for a specific shard.
    pub fn get_shard_root(&self, shard_id: ShardId) -> Option<Hash256> {
        self.shard_roots
            .iter()
            .find(|(id, _)| *id == shard_id)
            .map(|(_, root)| *root)
    }

    /// Checks if proof includes a specific shard.
    pub fn includes_shard(&self, shard_id: ShardId) -> bool {
        self.shard_roots.iter().any(|(id, _)| *id == shard_id)
    }
}

/// Aggregates proofs from multiple shards.
pub struct ProofAggregator {
    /// Number of shards.
    num_shards: u16,
    /// Cache of recent aggregated proofs.
    cache: Vec<AggregatedProof>,
    /// Maximum cache size.
    max_cache_size: usize,
}

impl ProofAggregator {
    /// Creates a new proof aggregator.
    pub fn new(num_shards: u16) -> Self {
        Self {
            num_shards,
            cache: Vec::new(),
            max_cache_size: 100,
        }
    }

    /// Aggregates proofs from the specified shards.
    pub fn aggregate(
        &self,
        shard_ids: &[ShardId],
        state_manager: &ShardStateManager,
    ) -> ShardResult<AggregatedProof> {
        // Validate shard IDs
        for &id in shard_ids {
            if id >= self.num_shards {
                return Err(ShardError::InvalidShardId(id, self.num_shards - 1));
            }
        }

        // Collect state roots
        let mut shard_roots = Vec::new();
        for &shard_id in shard_ids {
            if let Some(root) = state_manager.get_state_root(shard_id) {
                shard_roots.push((shard_id, root));
            } else {
                return Err(ShardError::ShardNotFound(shard_id));
            }
        }

        // Sort by shard ID for determinism
        shard_roots.sort_by_key(|(id, _)| *id);

        // Compute combined root
        let combined_root = AggregatedProof::compute_combined_root(&shard_roots);

        // Generate beacon block hash (in production, from actual beacon chain)
        let mut hasher = blake3::Hasher::new();
        hasher.update(combined_root.as_bytes());
        hasher.update(&(std::time::UNIX_EPOCH.elapsed().unwrap().as_secs()).to_le_bytes());
        let beacon_block = Hash256::from_bytes(*hasher.finalize().as_bytes());

        Ok(AggregatedProof {
            beacon_block,
            shard_roots,
            combined_root,
            proofs: Vec::new(),
            timestamp: std::time::UNIX_EPOCH.elapsed().unwrap().as_secs(),
        })
    }

    /// Aggregates proofs from all shards.
    pub fn aggregate_all(
        &self,
        state_manager: &ShardStateManager,
    ) -> ShardResult<AggregatedProof> {
        let all_ids: Vec<ShardId> = (0..self.num_shards).collect();
        self.aggregate(&all_ids, state_manager)
    }

    /// Caches an aggregated proof.
    pub fn cache_proof(&mut self, proof: AggregatedProof) {
        if self.cache.len() >= self.max_cache_size {
            self.cache.remove(0);
        }
        self.cache.push(proof);
    }

    /// Gets a cached proof by beacon block.
    pub fn get_cached(&self, beacon_block: &Hash256) -> Option<&AggregatedProof> {
        self.cache.iter().find(|p| &p.beacon_block == beacon_block)
    }

    /// Gets the most recent cached proof.
    pub fn latest_proof(&self) -> Option<&AggregatedProof> {
        self.cache.last()
    }

    /// Verifies a proof against the current state.
    pub fn verify_against_state(
        &self,
        proof: &AggregatedProof,
        state_manager: &ShardStateManager,
    ) -> bool {
        // First verify internal consistency
        if !proof.verify() {
            return false;
        }

        // Verify each shard root matches current state
        for (shard_id, proof_root) in &proof.shard_roots {
            if let Some(current_root) = state_manager.get_state_root(*shard_id) {
                if current_root != *proof_root {
                    return false;
                }
            } else {
                return false;
            }
        }

        true
    }

    /// Clears the cache.
    pub fn clear_cache(&mut self) {
        self.cache.clear();
    }

    /// Returns cache statistics.
    pub fn cache_stats(&self) -> CacheStats {
        CacheStats {
            cached_proofs: self.cache.len(),
            max_size: self.max_cache_size,
            oldest_timestamp: self.cache.first().map(|p| p.timestamp),
            newest_timestamp: self.cache.last().map(|p| p.timestamp),
        }
    }
}

/// Cache statistics.
#[derive(Clone, Debug)]
pub struct CacheStats {
    /// Number of cached proofs.
    pub cached_proofs: usize,
    /// Maximum cache size.
    pub max_size: usize,
    /// Oldest proof timestamp.
    pub oldest_timestamp: Option<u64>,
    /// Newest proof timestamp.
    pub newest_timestamp: Option<u64>,
}

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

    #[test]
    fn test_combined_root_deterministic() {
        let roots = vec![
            (0, Hash256::from_bytes([1u8; 32])),
            (1, Hash256::from_bytes([2u8; 32])),
        ];

        let root1 = AggregatedProof::compute_combined_root(&roots);
        let root2 = AggregatedProof::compute_combined_root(&roots);

        assert_eq!(root1, root2);
    }

    #[test]
    fn test_combined_root_order_matters() {
        let roots1 = vec![
            (0, Hash256::from_bytes([1u8; 32])),
            (1, Hash256::from_bytes([2u8; 32])),
        ];
        let roots2 = vec![
            (1, Hash256::from_bytes([2u8; 32])),
            (0, Hash256::from_bytes([1u8; 32])),
        ];

        let root1 = AggregatedProof::compute_combined_root(&roots1);
        let root2 = AggregatedProof::compute_combined_root(&roots2);

        // Different order = different root
        assert_ne!(root1, root2);
    }

    #[test]
    fn test_aggregate_proofs() {
        let state_manager = ShardStateManager::new(4);
        let aggregator = ProofAggregator::new(4);

        let proof = aggregator.aggregate(&[0, 1, 2], &state_manager).unwrap();

        assert_eq!(proof.shard_roots.len(), 3);
        assert!(proof.verify());
    }

    #[test]
    fn test_aggregate_all() {
        let state_manager = ShardStateManager::new(4);
        let aggregator = ProofAggregator::new(4);

        let proof = aggregator.aggregate_all(&state_manager).unwrap();

        assert_eq!(proof.shard_roots.len(), 4);
        assert!(proof.verify());
    }

    #[test]
    fn test_proof_verification() {
        let state_manager = ShardStateManager::new(4);
        let aggregator = ProofAggregator::new(4);

        let proof = aggregator.aggregate_all(&state_manager).unwrap();

        // Should verify against current state
        assert!(aggregator.verify_against_state(&proof, &state_manager));
    }

    #[test]
    fn test_cache() {
        let state_manager = ShardStateManager::new(4);
        let mut aggregator = ProofAggregator::new(4);

        let proof = aggregator.aggregate_all(&state_manager).unwrap();
        let beacon = proof.beacon_block;

        aggregator.cache_proof(proof);

        assert!(aggregator.get_cached(&beacon).is_some());
        assert!(aggregator.latest_proof().is_some());
    }

    #[test]
    fn test_invalid_shard_error() {
        let state_manager = ShardStateManager::new(4);
        let aggregator = ProofAggregator::new(4);

        let result = aggregator.aggregate(&[0, 10], &state_manager);
        assert!(result.is_err());
    }
}
