//! Device Mesh abstraction for multi-dimensional device organization.
//!
//! A DeviceMesh organizes devices (GPUs) into a logical multi-dimensional grid,
//! enabling hybrid parallelism strategies. Each dimension of the mesh can be
//! used for a different parallelism strategy (data, tensor, pipeline, etc.).
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_distributed::device_mesh::DeviceMesh;
//!
//! // Create a 2D mesh: 2 nodes x 4 GPUs = 8 total
//! // Dimension 0 ("dp") for data parallelism
//! // Dimension 1 ("tp") for tensor parallelism
//! let mesh = DeviceMesh::new(
//!     vec![2, 4],
//!     vec!["dp".to_string(), "tp".to_string()],
//! )?;
//!
//! // Get the process group for tensor parallelism
//! let tp_group = mesh.get_group("tp")?;
//! ```

use crate::error::{DistributedError, Result};
use crate::group::ProcessGroup;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

// =============================================================================
// Mesh Dimension
// =============================================================================

/// Named dimension of the device mesh.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeshDimension {
    /// Name of this dimension (e.g., "dp", "tp", "pp")
    pub name: String,
    /// Size of this dimension (number of devices)
    pub size: usize,
    /// Index in the mesh shape
    pub index: usize,
}

impl MeshDimension {
    /// Create a new mesh dimension.
    pub fn new(name: impl Into<String>, size: usize, index: usize) -> Self {
        Self {
            name: name.into(),
            size,
            index,
        }
    }
}

// =============================================================================
// Device Mesh
// =============================================================================

/// Multi-dimensional device mesh for organizing distributed computation.
///
/// The mesh provides a logical view of devices and enables creating sub-groups
/// for different parallelism dimensions.
#[derive(Debug)]
pub struct DeviceMesh {
    /// Shape of the mesh (e.g., [2, 4] for 2x4 mesh)
    shape: Vec<usize>,

    /// Named dimensions
    dimensions: Vec<MeshDimension>,

    /// Name to dimension index mapping
    name_to_dim: HashMap<String, usize>,

    /// Total number of devices
    world_size: usize,

    /// Local rank in the global mesh
    local_rank: AtomicUsize,

    /// Device list (flattened)
    devices: Vec<DeviceInfo>,

    /// Process groups for each dimension (lazily created)
    dimension_groups: RwLock<HashMap<usize, Arc<ProcessGroup>>>,

    /// Optional parent mesh (for submesh creation)
    parent: Option<Arc<DeviceMesh>>,
}

/// Information about a device in the mesh.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceInfo {
    /// Global rank of this device
    pub global_rank: usize,
    /// Coordinate in the mesh
    pub coordinate: Vec<usize>,
    /// Device type (e.g., "cuda:0", "cpu")
    pub device_type: String,
    /// Node this device belongs to
    pub node_id: usize,
    /// Local device index on the node
    pub local_device_id: usize,
}

impl DeviceMesh {
    /// Create a new device mesh with the given shape and dimension names.
    ///
    /// # Arguments
    /// * `shape` - Shape of the mesh (e.g., [2, 4] for 2x4)
    /// * `dim_names` - Names for each dimension (e.g., ["dp", "tp"])
    ///
    /// # Returns
    /// A new DeviceMesh
    pub fn new(shape: Vec<usize>, dim_names: Vec<impl Into<String>>) -> Result<Self> {
        if shape.len() != dim_names.len() {
            return Err(DistributedError::configuration(format!(
                "Shape length {} must match dim_names length {}",
                shape.len(),
                dim_names.len()
            )));
        }

        let world_size: usize = shape.iter().product();
        if world_size == 0 {
            return Err(DistributedError::configuration(
                "Mesh cannot have zero devices",
            ));
        }

        let dimensions: Vec<MeshDimension> = dim_names
            .into_iter()
            .enumerate()
            .map(|(i, name)| MeshDimension::new(name, shape[i], i))
            .collect();

        let name_to_dim: HashMap<String, usize> = dimensions
            .iter()
            .map(|d| (d.name.clone(), d.index))
            .collect();

        // Create device info for all devices
        let devices = Self::create_device_list(&shape, world_size);

        Ok(Self {
            shape,
            dimensions,
            name_to_dim,
            world_size,
            local_rank: AtomicUsize::new(0),
            devices,
            dimension_groups: RwLock::new(HashMap::new()),
            parent: None,
        })
    }

    /// Create a simple 1D mesh.
    pub fn new_simple(size: usize, dim_name: impl Into<String>) -> Self {
        Self::new(vec![size], vec![dim_name]).expect("Simple mesh creation should not fail")
    }

    /// Create device list from shape.
    fn create_device_list(shape: &[usize], world_size: usize) -> Vec<DeviceInfo> {
        let mut devices = Vec::with_capacity(world_size);

        for global_rank in 0..world_size {
            let coordinate = Self::rank_to_coordinate_static(global_rank, shape);
            devices.push(DeviceInfo {
                global_rank,
                coordinate,
                device_type: format!("cuda:{}", global_rank % 8), // Assume 8 GPUs per node
                node_id: global_rank / 8,
                local_device_id: global_rank % 8,
            });
        }

        devices
    }

    /// Convert global rank to mesh coordinate.
    fn rank_to_coordinate_static(rank: usize, shape: &[usize]) -> Vec<usize> {
        let mut coord = vec![0; shape.len()];
        let mut remaining = rank;

        for i in (0..shape.len()).rev() {
            coord[i] = remaining % shape[i];
            remaining /= shape[i];
        }

        coord
    }

    /// Get the shape of the mesh.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// Get the number of dimensions.
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    /// Get total number of devices.
    pub fn world_size(&self) -> usize {
        self.world_size
    }

    /// Get the local rank.
    pub fn local_rank(&self) -> usize {
        self.local_rank.load(Ordering::SeqCst)
    }

    /// Set the local rank.
    pub fn set_local_rank(&self, rank: usize) {
        self.local_rank.store(rank, Ordering::SeqCst);
    }

    /// Get dimension by name.
    pub fn get_dimension(&self, name: &str) -> Option<&MeshDimension> {
        self.name_to_dim
            .get(name)
            .and_then(|&idx| self.dimensions.get(idx))
    }

    /// Get dimension by index.
    pub fn get_dimension_by_index(&self, index: usize) -> Option<&MeshDimension> {
        self.dimensions.get(index)
    }

    /// Get the size of a dimension by name.
    pub fn dim_size(&self, name: &str) -> Option<usize> {
        self.get_dimension(name).map(|d| d.size)
    }

    /// Convert global rank to mesh coordinate.
    pub fn rank_to_coordinate(&self, rank: usize) -> Vec<usize> {
        Self::rank_to_coordinate_static(rank, &self.shape)
    }

    /// Convert mesh coordinate to global rank.
    pub fn coordinate_to_rank(&self, coord: &[usize]) -> Result<usize> {
        if coord.len() != self.shape.len() {
            return Err(DistributedError::configuration(format!(
                "Coordinate length {} must match mesh dimensions {}",
                coord.len(),
                self.shape.len()
            )));
        }

        for (i, (&c, &s)) in coord.iter().zip(self.shape.iter()).enumerate() {
            if c >= s {
                return Err(DistributedError::configuration(format!(
                    "Coordinate[{}] = {} exceeds dimension size {}",
                    i, c, s
                )));
            }
        }

        let mut rank = 0;
        let mut multiplier = 1;

        for i in (0..self.shape.len()).rev() {
            rank += coord[i] * multiplier;
            multiplier *= self.shape[i];
        }

        Ok(rank)
    }

    /// Get the local coordinate of this process in the mesh.
    pub fn local_coordinate(&self) -> Vec<usize> {
        self.rank_to_coordinate(self.local_rank())
    }

    /// Get device info for a rank.
    pub fn get_device(&self, rank: usize) -> Option<&DeviceInfo> {
        self.devices.get(rank)
    }

    /// Get all devices.
    pub fn devices(&self) -> &[DeviceInfo] {
        &self.devices
    }

    // =========================================================================
    // Process Group Operations
    // =========================================================================

    /// Get ranks along a mesh dimension for the current coordinate.
    ///
    /// Returns all ranks that share the same coordinates on all other dimensions.
    pub fn get_ranks_along_dim(&self, dim: usize) -> Result<Vec<usize>> {
        if dim >= self.ndim() {
            return Err(DistributedError::configuration(format!(
                "Dimension {} out of range for {}D mesh",
                dim,
                self.ndim()
            )));
        }

        let local_coord = self.local_coordinate();
        let mut ranks = Vec::with_capacity(self.shape[dim]);

        for i in 0..self.shape[dim] {
            let mut coord = local_coord.clone();
            coord[dim] = i;
            ranks.push(self.coordinate_to_rank(&coord)?);
        }

        Ok(ranks)
    }

    /// Get ranks along a mesh dimension by name.
    pub fn get_ranks_along_dim_by_name(&self, name: &str) -> Result<Vec<usize>> {
        let dim = self.name_to_dim.get(name).ok_or_else(|| {
            DistributedError::configuration(format!("Unknown dimension: {}", name))
        })?;
        self.get_ranks_along_dim(*dim)
    }

    /// Get the rank of this process within a dimension.
    pub fn get_rank_in_dim(&self, dim: usize) -> usize {
        self.local_coordinate().get(dim).copied().unwrap_or(0)
    }

    /// Get the rank of this process within a dimension by name.
    pub fn get_rank_in_dim_by_name(&self, name: &str) -> Option<usize> {
        self.name_to_dim
            .get(name)
            .map(|&dim| self.get_rank_in_dim(dim))
    }

    /// Get or create a process group for a dimension.
    ///
    /// Creates a ProcessGroup containing all ranks along the specified dimension
    /// that share the same coordinates on all other dimensions. The group uses
    /// RNCCL backend when the feature is enabled.
    pub fn get_process_group(&self, dim: usize) -> Result<Arc<ProcessGroup>> {
        // Check cache first
        {
            let groups = self.dimension_groups.read();
            if let Some(pg) = groups.get(&dim) {
                return Ok(pg.clone());
            }
        }

        // Create new process group
        let ranks = self.get_ranks_along_dim(dim)?;
        let local_rank = self.local_rank();

        // Find our position in the dimension
        let rank_in_dim = ranks.iter().position(|&r| r == local_rank).unwrap_or(0);
        let dim_size = ranks.len();

        // Create process group with RNCCL backend when available
        #[cfg(feature = "rnccl")]
        let pg = {
            use crate::backend::Backend;
            let world_info =
                crate::group::WorldInfo::new(dim_size as i32, rank_in_dim as i32, Backend::Rnccl);
            ProcessGroup::new(Backend::Rnccl, world_info)?
        };

        #[cfg(not(feature = "rnccl"))]
        let pg = {
            use crate::backend::Backend;
            let world_info =
                crate::group::WorldInfo::new(dim_size as i32, rank_in_dim as i32, Backend::Cpu);
            ProcessGroup::new(Backend::Cpu, world_info)?
        };

        let pg = Arc::new(pg);

        // Cache and return
        {
            let mut groups = self.dimension_groups.write();
            groups.insert(dim, pg.clone());
        }

        Ok(pg)
    }

    /// Get or create a process group for a dimension by name.
    pub fn get_process_group_by_name(&self, name: &str) -> Result<Arc<ProcessGroup>> {
        let dim = self.name_to_dim.get(name).ok_or_else(|| {
            DistributedError::configuration(format!("Unknown dimension: {}", name))
        })?;
        self.get_process_group(*dim)
    }

    // =========================================================================
    // Submesh Operations
    // =========================================================================

    /// Create a submesh by slicing along specified dimensions.
    ///
    /// # Arguments
    /// * `dim_indices` - Which dimensions to include in the submesh
    ///
    /// # Returns
    /// A new DeviceMesh representing the submesh
    pub fn submesh(&self, dim_names: &[&str]) -> Result<Arc<DeviceMesh>> {
        let dim_indices: Result<Vec<usize>> = dim_names
            .iter()
            .map(|name| {
                self.name_to_dim.get(*name).copied().ok_or_else(|| {
                    DistributedError::configuration(format!("Unknown dimension: {}", name))
                })
            })
            .collect();
        let dim_indices = dim_indices?;

        let new_shape: Vec<usize> = dim_indices.iter().map(|&i| self.shape[i]).collect();
        let new_dim_names: Vec<String> = dim_indices
            .iter()
            .map(|&i| self.dimensions[i].name.clone())
            .collect();

        let submesh = DeviceMesh {
            shape: new_shape.clone(),
            dimensions: new_dim_names
                .iter()
                .enumerate()
                .map(|(i, name)| MeshDimension::new(name.clone(), new_shape[i], i))
                .collect(),
            name_to_dim: new_dim_names
                .iter()
                .enumerate()
                .map(|(i, name)| (name.clone(), i))
                .collect(),
            world_size: new_shape.iter().product(),
            local_rank: AtomicUsize::new(self.local_rank()),
            devices: Vec::new(), // Submesh devices would be computed
            dimension_groups: RwLock::new(HashMap::new()),
            parent: Some(Arc::new(DeviceMesh {
                shape: self.shape.clone(),
                dimensions: self.dimensions.clone(),
                name_to_dim: self.name_to_dim.clone(),
                world_size: self.world_size,
                local_rank: AtomicUsize::new(self.local_rank()),
                devices: self.devices.clone(),
                dimension_groups: RwLock::new(HashMap::new()),
                parent: None,
            })),
        };

        Ok(Arc::new(submesh))
    }

    /// Flatten the mesh to 1D.
    pub fn flatten(&self) -> Arc<DeviceMesh> {
        Arc::new(DeviceMesh::new_simple(self.world_size, "flat"))
    }

    // =========================================================================
    // Collective Operations
    // =========================================================================

    /// Check if two ranks are neighbors along a dimension.
    pub fn are_neighbors(&self, rank1: usize, rank2: usize, dim: usize) -> bool {
        let coord1 = self.rank_to_coordinate(rank1);
        let coord2 = self.rank_to_coordinate(rank2);

        // Must differ only in the specified dimension
        for (i, (&c1, &c2)) in coord1.iter().zip(coord2.iter()).enumerate() {
            if i == dim {
                if (c1 as isize - c2 as isize).abs() != 1 {
                    return false;
                }
            } else if c1 != c2 {
                return false;
            }
        }

        true
    }

    /// Get neighbor ranks along a dimension.
    pub fn get_neighbors(&self, rank: usize, dim: usize) -> (Option<usize>, Option<usize>) {
        let coord = self.rank_to_coordinate(rank);
        let pos = coord[dim];

        let prev = if pos > 0 {
            let mut prev_coord = coord.clone();
            prev_coord[dim] = pos - 1;
            self.coordinate_to_rank(&prev_coord).ok()
        } else {
            None
        };

        let next = if pos + 1 < self.shape[dim] {
            let mut next_coord = coord.clone();
            next_coord[dim] = pos + 1;
            self.coordinate_to_rank(&next_coord).ok()
        } else {
            None
        };

        (prev, next)
    }

    /// Get the distance between two ranks (number of hops).
    pub fn distance(&self, rank1: usize, rank2: usize) -> usize {
        let coord1 = self.rank_to_coordinate(rank1);
        let coord2 = self.rank_to_coordinate(rank2);

        coord1
            .iter()
            .zip(coord2.iter())
            .map(|(a, b)| (*a as isize - *b as isize).unsigned_abs())
            .sum()
    }
}

// =============================================================================
// Mesh Builder
// =============================================================================

/// Builder for creating device meshes with fluent API.
#[derive(Default)]
pub struct DeviceMeshBuilder {
    shape: Vec<usize>,
    dim_names: Vec<String>,
    local_rank: Option<usize>,
}

impl DeviceMeshBuilder {
    /// Create a new builder.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a dimension to the mesh.
    pub fn add_dim(mut self, name: impl Into<String>, size: usize) -> Self {
        self.dim_names.push(name.into());
        self.shape.push(size);
        self
    }

    /// Set the local rank.
    pub fn with_local_rank(mut self, rank: usize) -> Self {
        self.local_rank = Some(rank);
        self
    }

    /// Build the device mesh.
    pub fn build(self) -> Result<DeviceMesh> {
        let mesh = DeviceMesh::new(self.shape, self.dim_names)?;
        if let Some(rank) = self.local_rank {
            mesh.set_local_rank(rank);
        }
        Ok(mesh)
    }
}

// =============================================================================
// Common Mesh Patterns
// =============================================================================

impl DeviceMesh {
    /// Create a 2D mesh for data + tensor parallelism.
    pub fn data_tensor_parallel(dp_size: usize, tp_size: usize) -> Result<Self> {
        DeviceMeshBuilder::new()
            .add_dim("dp", dp_size)
            .add_dim("tp", tp_size)
            .build()
    }

    /// Create a 3D mesh for data + tensor + pipeline parallelism.
    pub fn parallelism_3d(dp_size: usize, tp_size: usize, pp_size: usize) -> Result<Self> {
        DeviceMeshBuilder::new()
            .add_dim("dp", dp_size)
            .add_dim("tp", tp_size)
            .add_dim("pp", pp_size)
            .build()
    }

    /// Create a mesh for expert parallelism (MoE).
    pub fn expert_parallel(dp_size: usize, ep_size: usize, tp_size: usize) -> Result<Self> {
        DeviceMeshBuilder::new()
            .add_dim("dp", dp_size)
            .add_dim("ep", ep_size)
            .add_dim("tp", tp_size)
            .build()
    }
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_mesh_creation() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
        assert_eq!(mesh.shape(), &[2, 4]);
        assert_eq!(mesh.world_size(), 8);
        assert_eq!(mesh.ndim(), 2);
    }

    #[test]
    fn test_simple_mesh() {
        let mesh = DeviceMesh::new_simple(4, "dp");
        assert_eq!(mesh.shape(), &[4]);
        assert_eq!(mesh.world_size(), 4);
    }

    #[test]
    fn test_rank_to_coordinate() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();

        assert_eq!(mesh.rank_to_coordinate(0), vec![0, 0]);
        assert_eq!(mesh.rank_to_coordinate(1), vec![0, 1]);
        assert_eq!(mesh.rank_to_coordinate(4), vec![1, 0]);
        assert_eq!(mesh.rank_to_coordinate(7), vec![1, 3]);
    }

    #[test]
    fn test_coordinate_to_rank() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();

        assert_eq!(mesh.coordinate_to_rank(&[0, 0]).unwrap(), 0);
        assert_eq!(mesh.coordinate_to_rank(&[0, 1]).unwrap(), 1);
        assert_eq!(mesh.coordinate_to_rank(&[1, 0]).unwrap(), 4);
        assert_eq!(mesh.coordinate_to_rank(&[1, 3]).unwrap(), 7);
    }

    #[test]
    fn test_get_ranks_along_dim() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
        mesh.set_local_rank(5); // Coordinate [1, 1]

        // Along dp dimension (dim 0): ranks with tp=1
        let dp_ranks = mesh.get_ranks_along_dim(0).unwrap();
        assert_eq!(dp_ranks, vec![1, 5]);

        // Along tp dimension (dim 1): ranks with dp=1
        let tp_ranks = mesh.get_ranks_along_dim(1).unwrap();
        assert_eq!(tp_ranks, vec![4, 5, 6, 7]);
    }

    #[test]
    fn test_dim_size_by_name() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();
        assert_eq!(mesh.dim_size("dp"), Some(2));
        assert_eq!(mesh.dim_size("tp"), Some(4));
        assert_eq!(mesh.dim_size("pp"), None);
    }

    #[test]
    fn test_neighbors() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();

        // Rank 5 is at [1, 1]
        let (prev, next) = mesh.get_neighbors(5, 1); // along tp
        assert_eq!(prev, Some(4)); // [1, 0]
        assert_eq!(next, Some(6)); // [1, 2]

        // Edge cases
        let (prev, next) = mesh.get_neighbors(4, 1); // [1, 0]
        assert_eq!(prev, None);
        assert_eq!(next, Some(5));

        let (prev, next) = mesh.get_neighbors(7, 1); // [1, 3]
        assert_eq!(prev, Some(6));
        assert_eq!(next, None);
    }

    #[test]
    fn test_distance() {
        let mesh = DeviceMesh::new(vec![2, 4], vec!["dp", "tp"]).unwrap();

        assert_eq!(mesh.distance(0, 0), 0);
        assert_eq!(mesh.distance(0, 1), 1); // [0,0] to [0,1]
        assert_eq!(mesh.distance(0, 4), 1); // [0,0] to [1,0]
        assert_eq!(mesh.distance(0, 7), 4); // [0,0] to [1,3]
    }

    #[test]
    fn test_builder() {
        let mesh = DeviceMeshBuilder::new()
            .add_dim("dp", 2)
            .add_dim("tp", 4)
            .with_local_rank(5)
            .build()
            .unwrap();

        assert_eq!(mesh.shape(), &[2, 4]);
        assert_eq!(mesh.local_rank(), 5);
    }

    #[test]
    fn test_3d_parallelism() {
        let mesh = DeviceMesh::parallelism_3d(2, 2, 2).unwrap();
        assert_eq!(mesh.shape(), &[2, 2, 2]);
        assert_eq!(mesh.world_size(), 8);
        assert!(mesh.get_dimension("dp").is_some());
        assert!(mesh.get_dimension("tp").is_some());
        assert!(mesh.get_dimension("pp").is_some());
    }
}