//! Process group management for distributed training
//!
//! This module provides the core ProcessGroup abstraction that manages
//! distributed process communication, including initialization, splitting,
//! merging, and cleanup operations.

use crate::backend::{Backend, BackendConfig, BackendImpl};
#[cfg(any(feature = "rnccl", feature = "nccl"))]
use crate::comm::CommunicationPrimitive;
use crate::error::{DistributedError, Result};
#[cfg(feature = "nccl")]
use crate::nccl::NcclCommunicator;
#[cfg(feature = "nccl")]
use crate::nccl_comm::NcclCommunicationPrimitive;
#[cfg(feature = "rnccl")]
use crate::rnccl_backend::RncclBackend;
use crate::{Tensor, TensorShape};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use uuid::Uuid;

/// Information about the distributed world
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorldInfo {
    /// Rank of this process (0-indexed)
    pub rank: usize,
    /// Total number of processes in the world
    pub world_size: usize,
    /// Local rank within the node
    pub local_rank: usize,
    /// Number of processes on local node
    pub local_world_size: usize,
    /// Master address for rendezvous
    pub master_addr: String,
    /// Master port for rendezvous
    pub master_port: u16,
    /// Unique identifier for this process group
    pub group_id: Option<Uuid>,
    /// Backend being used (derived from ProcessGroup)
    #[serde(skip)]
    pub backend: Option<Backend>,
    /// Timestamp when the group was created (seconds since UNIX epoch)
    pub created_at: u64,
    /// Instant when the group was created (not serialized)
    #[serde(skip)]
    pub created_instant: Option<Instant>,
}

impl WorldInfo {
    /// Create new world info for simple cases
    pub fn new(world_size: i32, rank: i32, backend: Backend) -> Self {
        let now = Instant::now();
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        Self {
            rank: rank as usize,
            world_size: world_size as usize,
            local_rank: rank as usize,
            local_world_size: world_size as usize,
            master_addr: "127.0.0.1".to_string(),
            master_port: 29500,
            group_id: Some(Uuid::new_v4()),
            backend: Some(backend),
            created_at: timestamp,
            created_instant: Some(now),
        }
    }

    /// Check if this process is the root (rank 0)
    pub fn is_root(&self) -> bool {
        self.rank == 0
    }

    /// Get elapsed time since group creation
    pub fn elapsed(&self) -> std::time::Duration {
        self.created_instant.unwrap_or_else(Instant::now).elapsed()
    }
}

/// Internal state of a process group
#[derive(Debug)]
struct ProcessGroupState {
    world_info: WorldInfo,
    config: BackendConfig,
    backend_impl: Option<BackendImpl>,
    is_initialized: bool,
    is_cleaned_up: bool,
}

impl ProcessGroupState {
    fn new(world_size: i32, rank: i32, backend: Backend, config: BackendConfig) -> Self {
        Self {
            world_info: WorldInfo::new(world_size, rank, backend),
            config,
            backend_impl: None,
            is_initialized: false,
            is_cleaned_up: false,
        }
    }
}

/// Process group for distributed training
///
/// A ProcessGroup represents a collection of processes that can communicate
/// with each other using collective operations. It manages the lifetime of
/// the communication backend and provides a safe, high-level interface.
#[derive(Debug, Clone)]
pub struct ProcessGroup {
    state: Arc<RwLock<ProcessGroupState>>,
}

impl ProcessGroup {
    /// Create a new process group with WorldInfo
    pub fn new(backend: Backend, world_info: WorldInfo) -> Result<Self> {
        // Create default configuration for the backend
        let config = match backend {
            Backend::Nccl => BackendConfig::nccl(),
            Backend::Rccl => BackendConfig::rccl(),
            Backend::Rnccl => BackendConfig::nccl(), // Rust NCCL uses same config as NCCL
            Backend::Mpi => BackendConfig::mpi(),
            Backend::Cpu => BackendConfig::cpu(),
            Backend::Tcp => BackendConfig::tcp(),
        };

        // Update world_info with backend
        let mut updated_world_info = world_info;
        updated_world_info.backend = Some(backend);
        if updated_world_info.created_instant.is_none() {
            updated_world_info.created_instant = Some(Instant::now());
        }

        let state = ProcessGroupState {
            world_info: updated_world_info,
            config,
            backend_impl: None,   // Backend implementation is initialized lazily
            is_initialized: true, // For simplicity, mark as initialized
            is_cleaned_up: false,
        };

        Ok(Self {
            state: Arc::new(RwLock::new(state)),
        })
    }

    /// Create a new process group (legacy interface)
    ///
    /// # Arguments
    /// * `backend` - Communication backend to use
    /// * `world_size` - Total number of processes
    /// * `rank` - Rank of this process (0-indexed)
    /// * `config` - Backend-specific configuration
    ///
    /// # Returns
    /// * `Result<Self>` - New process group or error
    ///
    /// # Examples
    /// ```rust,ignore
    /// use rtx_distributed::{Backend, BackendConfig, ProcessGroup};
    ///
    /// let config = BackendConfig::nccl();
    /// let pg = ProcessGroup::new(Backend::Nccl, 4, 0, config).await?;
    /// ```
    pub async fn new_with_config(
        backend: Backend,
        world_size: i32,
        rank: i32,
        config: BackendConfig,
    ) -> Result<Self> {
        // Validate parameters
        if world_size <= 0 {
            return Err(DistributedError::process_group(
                "world_size must be positive",
            ));
        }

        if rank < 0 || rank >= world_size {
            return Err(DistributedError::process_group(format!(
                "rank {rank} must be in range [0, {world_size})"
            )));
        }

        // Validate configuration
        config.validate()?;

        if config.backend != backend {
            return Err(DistributedError::configuration(
                "backend mismatch between parameter and config",
            ));
        }

        // Create process group state
        let state = ProcessGroupState::new(world_size, rank, backend, config);
        let pg = Self {
            state: Arc::new(RwLock::new(state)),
        };

        // Initialize the backend
        pg.initialize_backend().await?;

        Ok(pg)
    }

    /// Initialize the communication backend
    async fn initialize_backend(&self) -> Result<()> {
        let mut state = self.state.write();

        if state.is_initialized {
            return Ok(());
        }

        match state.world_info.backend {
            Some(Backend::Nccl) => {
                self.init_nccl(&mut state).await?;
            }
            Some(Backend::Rccl) => {
                self.init_rccl(&mut state).await?;
            }
            Some(Backend::Rnccl) => {
                self.init_rnccl(&mut state).await?;
            }
            Some(Backend::Mpi) => {
                self.init_mpi(&mut state).await?;
            }
            Some(Backend::Cpu) => {
                self.init_cpu(&mut state).await?;
            }
            Some(Backend::Tcp) => {
                self.init_tcp(&mut state).await?;
            }
            None => {
                return Err(DistributedError::process_group("backend not specified"));
            }
        }

        state.is_initialized = true;
        Ok(())
    }

    /// Initialize NCCL backend
    async fn init_nccl(&self, state: &mut ProcessGroupState) -> Result<()> {
        tracing::info!(
            "Initializing NCCL backend for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        #[cfg(feature = "nccl")]
        {
            // Create backend implementation
            let backend_impl = BackendImpl::new(state.config.clone()).await?;

            // Generate NCCL unique ID (in practice, this would be coordinated across all processes)
            let nccl_id = NcclCommunicator::get_unique_id().map_err(|e| {
                DistributedError::communication("nccl", format!("failed to get NCCL ID: {e}"))
            })?;

            // Initialize world communicator
            if let BackendImpl::Nccl(nccl_backend) = &backend_impl {
                nccl_backend
                    .init_world(
                        state.world_info.world_size,
                        state.world_info.rank,
                        state.world_info.rank as i32,
                        nccl_id,
                    )
                    .await?;
            } else {
                return Err(DistributedError::configuration(
                    "expected NCCL backend implementation",
                ));
            }

            state.backend_impl = Some(backend_impl);

            tracing::info!(
                "NCCL backend successfully initialized for rank {} of {}",
                state.world_info.rank,
                state.world_info.world_size
            );
        }

        #[cfg(not(feature = "nccl"))]
        {
            return Err(DistributedError::configuration(
                "NCCL backend not available - compile with nccl feature",
            ));
        }

        // Check timeout
        if state.world_info.elapsed() > state.config.timeout {
            return Err(DistributedError::communication(
                "nccl",
                "initialization timeout",
            ));
        }

        Ok(())
    }

    /// Initialize RCCL backend
    async fn init_rccl(&self, state: &mut ProcessGroupState) -> Result<()> {
        tracing::info!(
            "Initializing RCCL backend for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        // Simulate RCCL initialization
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        if state.world_info.elapsed() > state.config.timeout {
            return Err(DistributedError::communication(
                "rccl",
                "initialization timeout",
            ));
        }

        Ok(())
    }

    /// Initialize RNCCL (Rust-native NCCL) backend
    async fn init_rnccl(&self, state: &mut ProcessGroupState) -> Result<()> {
        tracing::info!(
            "Initializing RNCCL backend for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        #[cfg(feature = "rnccl")]
        {
            // Create RNCCL backend with rank and world_size
            let rnccl_backend = RncclBackend::new(
                state.world_info.rank as i32,
                state.world_info.world_size as i32,
            )?;

            // Store in BackendImpl
            state.backend_impl = Some(BackendImpl::Rnccl(Arc::new(rnccl_backend)));

            tracing::info!(
                "RNCCL backend successfully initialized for rank {} of {}",
                state.world_info.rank,
                state.world_info.world_size
            );
        }

        #[cfg(not(feature = "rnccl"))]
        {
            return Err(DistributedError::configuration(
                "RNCCL backend not available - compile with rnccl feature",
            ));
        }

        // Check timeout
        if state.world_info.elapsed() > state.config.timeout {
            return Err(DistributedError::communication(
                "rnccl",
                "initialization timeout",
            ));
        }

        Ok(())
    }

    /// Initialize MPI backend
    async fn init_mpi(&self, _state: &mut ProcessGroupState) -> Result<()> {
        #[cfg(feature = "mpi")]
        {
            tracing::info!("Initializing MPI backend");
            // MPI initialization would go here
            Ok(())
        }

        #[cfg(not(feature = "mpi"))]
        {
            Err(DistributedError::configuration(
                "MPI backend not available - compile with mpi feature",
            ))
        }
    }

    /// Initialize CPU backend (for testing)
    async fn init_cpu(&self, state: &mut ProcessGroupState) -> Result<()> {
        tracing::info!(
            "Initializing CPU backend for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        // Create CPU backend implementation
        let backend_impl = BackendImpl::new(state.config.clone()).await?;
        state.backend_impl = Some(backend_impl);

        // CPU backend is always available and fast to initialize
        tokio::time::sleep(std::time::Duration::from_millis(1)).await;

        tracing::info!(
            "CPU backend successfully initialized for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        Ok(())
    }

    /// Initialize TCP fallback backend
    async fn init_tcp(&self, state: &mut ProcessGroupState) -> Result<()> {
        tracing::info!(
            "Initializing TCP fallback backend for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        // Create TCP backend implementation
        let backend_impl = BackendImpl::new(state.config.clone()).await?;
        state.backend_impl = Some(backend_impl);

        tracing::info!(
            "TCP fallback backend successfully initialized for rank {} of {}",
            state.world_info.rank,
            state.world_info.world_size
        );

        Ok(())
    }

    /// Get world size
    pub fn world_size(&self) -> usize {
        self.state.read().world_info.world_size
    }

    /// Get rank
    pub fn rank(&self) -> usize {
        self.state.read().world_info.rank
    }

    /// Get backend
    pub fn backend(&self) -> Backend {
        self.state.read().world_info.backend.unwrap_or(Backend::Cpu)
    }

    /// Get world info
    pub fn world_info(&self) -> WorldInfo {
        self.state.read().world_info.clone()
    }

    /// Get communication primitive for this process group
    #[cfg(feature = "nccl")]
    pub fn communication_primitive(&self) -> Result<Option<NcclCommunicationPrimitive>> {
        let state = self.state.read();

        if let Some(BackendImpl::Nccl(nccl_backend)) = &state.backend_impl {
            Ok(Some(NcclCommunicationPrimitive::new(
                nccl_backend.clone(),
                state.world_info.world_size,
                state.world_info.rank,
            )))
        } else {
            Ok(None)
        }
    }

    /// Check if NCCL backend is available and initialized
    #[cfg(feature = "nccl")]
    pub fn has_nccl_backend(&self) -> bool {
        let state = self.state.read();
        matches!(&state.backend_impl, Some(BackendImpl::Nccl(_)))
    }

    /// Check if RNCCL backend is available and initialized
    #[cfg(feature = "rnccl")]
    pub fn has_rnccl_backend(&self) -> bool {
        let state = self.state.read();
        matches!(&state.backend_impl, Some(BackendImpl::Rnccl(_)))
    }

    /// Get the RNCCL backend reference for direct operations
    #[cfg(feature = "rnccl")]
    pub fn rnccl_backend(&self) -> Option<Arc<RncclBackend>> {
        let state = self.state.read();
        if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
            Some(backend.clone())
        } else {
            None
        }
    }

    /// AllReduce operation - sum tensors across all processes
    pub async fn all_reduce(&self, tensor: &mut Tensor, op: crate::comm::ReduceOp) -> Result<()> {
        use crate::comm::ReduceOp;

        // Check initialization
        {
            let state = self.state.read();
            if !state.is_initialized {
                return Err(DistributedError::communication(
                    "allreduce",
                    "process group not initialized",
                ));
            }
        }

        // Use NCCL backend if available
        #[cfg(feature = "nccl")]
        {
            let nccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Nccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(nccl_backend) = nccl_backend_opt {
                let comm_primitive =
                    NcclCommunicationPrimitive::new(nccl_backend, self.world_size(), self.rank());

                return comm_primitive.allreduce(tensor, op).await;
            }
        }

        // Use RNCCL backend if available
        #[cfg(feature = "rnccl")]
        {
            let rnccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(rnccl_backend) = rnccl_backend_opt {
                return rnccl_backend.allreduce(tensor, op).await;
            }
        }

        // Fallback to simulation for CPU backend
        match op {
            ReduceOp::Sum => {
                // In real implementation, this would sum across all ranks
                // For simulation, we modify the tensor to show it was processed
                let scale_factor = self.world_size() as f32;
                *tensor = tensor.mul_scalar(scale_factor)?;
            }
            ReduceOp::Max | ReduceOp::Min => {
                // For min/max, we don't need to scale
            }
            _ => {
                // Other ops would have their own implementations
            }
        }

        Ok(())
    }

    /// Broadcast tensor from root process to all others
    pub async fn broadcast(&self, tensor: &mut Tensor, root: usize) -> Result<()> {
        // Check initialization and bounds
        {
            let state = self.state.read();
            if !state.is_initialized {
                return Err(DistributedError::communication(
                    "broadcast",
                    "process group not initialized",
                ));
            }

            if root >= state.world_info.world_size {
                return Err(DistributedError::communication(
                    "broadcast",
                    format!(
                        "root rank {} out of bounds for world_size {}",
                        root, state.world_info.world_size
                    ),
                ));
            }
        }

        // Use NCCL backend if available
        #[cfg(feature = "nccl")]
        {
            let nccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Nccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(nccl_backend) = nccl_backend_opt {
                let comm_primitive =
                    NcclCommunicationPrimitive::new(nccl_backend, self.world_size(), self.rank());

                return comm_primitive.broadcast(tensor, root as i32).await;
            }
        }

        // Use RNCCL backend if available
        #[cfg(feature = "rnccl")]
        {
            let rnccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(rnccl_backend) = rnccl_backend_opt {
                return rnccl_backend.broadcast(tensor, root as i32).await;
            }
        }

        // Simulate broadcast - in real implementation, non-root processes would receive data
        tracing::debug!("Broadcasting tensor from rank {} to all processes", root);
        Ok(())
    }

    /// AllGather operation - gather tensors from all processes
    pub fn all_gather(&self, tensor: &Tensor) -> Result<Vec<Tensor>> {
        // Check initialization
        let world_size = {
            let state = self.state.read();
            if !state.is_initialized {
                return Err(DistributedError::communication(
                    "allgather",
                    "process group not initialized",
                ));
            }
            state.world_info.world_size
        };

        // Use RNCCL backend if available (blocking call to async fn)
        #[cfg(feature = "rnccl")]
        {
            let rnccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(rnccl_backend) = rnccl_backend_opt {
                use crate::comm::AllGatherOutput;
                // Use futures::executor for sync context
                let output = futures::executor::block_on(rnccl_backend.allgather(tensor))?;
                return match output {
                    AllGatherOutput::TensorList(tensors) => Ok(tensors),
                    AllGatherOutput::Tensor(t) => Ok(vec![t]),
                };
            }
        }

        // Simulate AllGather by creating copies for each rank
        let mut gathered = Vec::new();
        for _rank in 0..world_size {
            gathered.push(tensor.clone());
        }

        Ok(gathered)
    }

    /// ReduceScatter operation - reduce and scatter result
    pub fn reduce_scatter(&self, tensor: &Tensor, op: crate::comm::ReduceOp) -> Result<Tensor> {
        // Check initialization
        let (world_size, rank) = {
            let state = self.state.read();
            if !state.is_initialized {
                return Err(DistributedError::communication(
                    "reducescatter",
                    "process group not initialized",
                ));
            }
            (state.world_info.world_size, state.world_info.rank)
        };

        // Use RNCCL backend if available (blocking call to async fn)
        #[cfg(feature = "rnccl")]
        {
            let rnccl_backend_opt = {
                let state = self.state.read();
                if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
                    Some(backend.clone())
                } else {
                    None
                }
            };

            if let Some(rnccl_backend) = rnccl_backend_opt {
                return futures::executor::block_on(rnccl_backend.reduce_scatter(tensor, op));
            }
        }

        // Simulate ReduceScatter by returning a portion of the input tensor
        let elements_per_rank = tensor.numel() / world_size;
        let start_idx = rank * elements_per_rank;
        let _end_idx = start_idx + elements_per_rank;

        // Create a view/slice of the tensor (simplified)
        let result_shape = TensorShape::new(vec![elements_per_rank])?;
        let result = Tensor::zeros(result_shape, &crate::Device::default())?;

        Ok(result)
    }

    /// Split the process group
    pub async fn split(&self, _color: i32, key: i32) -> Result<Self> {
        let state = self.state.read();

        if !state.is_initialized {
            return Err(DistributedError::process_group(
                "cannot split uninitialized process group",
            ));
        }

        // For simplicity, create a new group with modified parameters
        // In a real implementation, this would use backend-specific splitting
        let new_world_size = state.world_info.world_size / 2; // Simplified
        let new_rank = (key as usize) % new_world_size;

        drop(state); // Release the lock before async call

        let new_config = self.state.read().config.clone();
        Self::new_with_config(
            self.backend(),
            new_world_size as i32,
            new_rank as i32,
            new_config,
        )
        .await
    }

    /// Merge multiple process groups
    pub async fn merge(groups: Vec<Self>) -> Result<Self> {
        if groups.is_empty() {
            return Err(DistributedError::process_group(
                "cannot merge empty list of groups",
            ));
        }

        let first_group = &groups[0];
        let backend = first_group.backend();
        let config = first_group.state.read().config.clone();

        // Calculate merged world size and rank
        let total_world_size: usize = groups.iter().map(ProcessGroup::world_size).sum();

        let new_rank = groups[0].rank(); // Simplified

        Self::new_with_config(backend, total_world_size as i32, new_rank as i32, config).await
    }

    /// Clean up the process group
    pub async fn cleanup(&self) -> Result<()> {
        let mut state = self.state.write();

        if state.is_cleaned_up {
            return Ok(());
        }

        tracing::info!(
            "Cleaning up process group for rank {}/{}",
            state.world_info.rank,
            state.world_info.world_size
        );

        // Clean up backend implementation
        if let Some(backend_impl) = state.backend_impl.take() {
            drop(state); // Release lock before async operation
            backend_impl.cleanup().await?;
            state = self.state.write(); // Re-acquire lock
        }

        // Additional backend-specific cleanup
        match state.world_info.backend {
            Some(Backend::Nccl | Backend::Rccl | Backend::Rnccl) => {
                // Additional GPU resource cleanup if needed
                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
            }
            Some(Backend::Mpi) => {
                #[cfg(feature = "mpi")]
                {
                    // MPI cleanup
                }
            }
            Some(Backend::Cpu | Backend::Tcp) | None => {
                // No special cleanup needed
            }
        }

        state.is_cleaned_up = true;
        tracing::info!(
            "Process group cleanup completed for rank {}/{}",
            state.world_info.rank,
            state.world_info.world_size
        );

        Ok(())
    }

    /// Synchronize all processes in the group (barrier operation)
    pub async fn barrier(&self) -> Result<()> {
        let state = self.state.read();

        if !state.is_initialized {
            return Err(DistributedError::process_group(
                "cannot perform barrier on uninitialized process group",
            ));
        }

        // For single process groups, barrier is a no-op
        if state.world_info.world_size == 1 {
            return Ok(());
        }

        // In a real implementation, this would use backend-specific barrier
        // For now, we'll just return success as a placeholder
        tracing::debug!(
            "Barrier operation for rank {} of {} (placeholder implementation)",
            state.world_info.rank,
            state.world_info.world_size
        );

        Ok(())
    }
}

// Automatic cleanup on drop
impl Drop for ProcessGroup {
    fn drop(&mut self) {
        let state = self.state.read();
        if state.is_initialized && !state.is_cleaned_up {
            // Log warning about unclean shutdown
            tracing::warn!(
                "ProcessGroup dropped without explicit cleanup for rank {} of {}",
                state.world_info.rank,
                state.world_info.world_size
            );
        }
    }
}

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

    #[tokio::test]
    async fn test_process_group_creation() {
        let config = BackendConfig::cpu();
        let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
            .await
            .unwrap();

        assert_eq!(pg.world_size(), 2);
        assert_eq!(pg.rank(), 0);
        assert_eq!(pg.backend(), Backend::Cpu);
    }

    #[tokio::test]
    async fn test_invalid_parameters() {
        let config = BackendConfig::cpu();

        // Invalid world size
        let result = ProcessGroup::new_with_config(Backend::Cpu, 0, 0, config.clone()).await;
        assert!(result.is_err());

        // Invalid rank
        let result = ProcessGroup::new_with_config(Backend::Cpu, 2, 2, config).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_world_info() {
        let config = BackendConfig::cpu();
        let pg = ProcessGroup::new_with_config(Backend::Cpu, 4, 1, config)
            .await
            .unwrap();

        let info = pg.world_info();
        assert_eq!(info.world_size, 4);
        assert_eq!(info.rank, 1);
        assert_eq!(info.backend, Some(Backend::Cpu));
        assert!(info.group_id.is_some());
        assert!(!info.is_root());
    }

    #[tokio::test]
    async fn test_cleanup() {
        let config = BackendConfig::cpu();
        let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
            .await
            .unwrap();

        assert!(pg.cleanup().await.is_ok());

        // Second cleanup should also work
        assert!(pg.cleanup().await.is_ok());
    }

    #[test]
    fn test_world_info_properties() {
        let info = WorldInfo::new(4, 0, Backend::Cpu);
        assert!(info.is_root());
        assert!(info.group_id.is_some());

        let info2 = WorldInfo::new(4, 1, Backend::Cpu);
        assert!(!info2.is_root());
    }
}