//! Sync operations for TestCluster.

use crate::error::{ChaosError, Result};
use std::collections::HashSet;
use tracing::{info, instrument};

use super::types::TestCluster;

/// Extension trait for sync operations.
pub trait TestClusterSyncExt {
    /// Syncs all nodes with each other (full mesh sync).
    ///
    /// Each alive node syncs from every other alive node.
    /// Respects network partitions.
    fn sync_all(&mut self) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Syncs a specific node from another specific node.
    fn sync_pair(
        &mut self,
        from: usize,
        to: usize,
    ) -> impl std::future::Future<Output = Result<usize>> + Send;
}

impl TestClusterSyncExt for TestCluster {
    #[instrument(skip(self))]
    async fn sync_all(&mut self) -> Result<()> {
        use super::access::TestClusterAccessExt;

        let alive_indices = self.alive_node_indices();
        info!(alive_count = alive_indices.len(), "Starting full mesh sync");

        // We need to be careful with borrowing - collect data first
        for i in 0..alive_indices.len() {
            for j in 0..alive_indices.len() {
                if i == j {
                    continue;
                }

                let from_idx = alive_indices[j];
                let to_idx = alive_indices[i];

                // Check if sync would be blocked
                if self.network.should_drop(from_idx, to_idx) {
                    continue;
                }

                // Collect data from source node
                let transfers: Vec<_> = {
                    let from_node = &self.nodes[from_idx];
                    from_node
                        .hash_to_data
                        .iter()
                        .map(|(h, (s, d))| (*h, s.clone(), d.clone()))
                        .collect()
                };

                // Apply latency
                self.network.apply_effects(from_idx, to_idx).await;

                // Merge into target node
                let to_node = &mut self.nodes[to_idx];
                let my_leaves: HashSet<_> = to_node.leaves().into_iter().collect();

                for (hash, subject, data) in transfers {
                    if !my_leaves.contains(&hash) {
                        let transfer =
                            stemedb_storage::crdt::AssertionTransfer { hash, data: data.clone() };

                        if to_node
                            .crdt_store
                            .merge_with_data(&subject, std::slice::from_ref(&transfer))
                            .await
                            .is_ok()
                        {
                            to_node.merkle_tree.insert(hash)?;
                            to_node.hash_to_data.insert(hash, (subject.clone(), data));
                        }
                    }
                }
            }
        }

        Ok(())
    }

    async fn sync_pair(&mut self, from: usize, to: usize) -> Result<usize> {
        if from >= self.nodes.len() {
            return Err(ChaosError::NodeNotFound(from));
        }
        if to >= self.nodes.len() {
            return Err(ChaosError::NodeNotFound(to));
        }
        if from == to {
            return Ok(0);
        }

        // Check partition
        if self.network.should_drop(from, to) {
            return Ok(0);
        }

        // Collect data from source
        let transfers: Vec<_> = {
            let from_node = &self.nodes[from];
            from_node.hash_to_data.iter().map(|(h, (s, d))| (*h, s.clone(), d.clone())).collect()
        };

        // Apply latency
        self.network.apply_effects(from, to).await;

        // Merge into target
        let to_node = &mut self.nodes[to];
        let my_leaves: HashSet<_> = to_node.leaves().into_iter().collect();
        let mut merged = 0;

        for (hash, subject, data) in transfers {
            if !my_leaves.contains(&hash) {
                let transfer =
                    stemedb_storage::crdt::AssertionTransfer { hash, data: data.clone() };

                if to_node
                    .crdt_store
                    .merge_with_data(&subject, std::slice::from_ref(&transfer))
                    .await
                    .is_ok()
                {
                    to_node.merkle_tree.insert(hash)?;
                    to_node.hash_to_data.insert(hash, (subject.clone(), data));
                    merged += 1;
                }
            }
        }

        Ok(merged)
    }
}
