class DeviceMeshParameters(BaseModel):
    """
    Configuration parameters for initializing Distributed Device Meshes.

    Attributes:
        pipeline_parallel: Degree of pipeline parallelism (PP).
        data_parallel_replicate: Degree of data parallel replication (DDP).
        data_parallel_shard: Degree of data parallel sharding (FSDP).
        context_parallel_replicate: Degree of context parallel (CP) replication.
        context_parallel_shard: Degree of context parallel (FSCP) sharding.
        tensor_parallel: Degree of tensor parallelism (TP).
        expert_parallel: Degree of expert parallelism (EP/MoE).
    """

    pipeline_parallel: int

    data_parallel_replicate: int
    data_parallel_shard: int

    context_parallel_replicate: int
    context_parallel_shard: int

    tensor_parallel: int

    expert_parallel: int

    @property
    def has_pipeline_parallel(self) -> bool:
        """Checks if pipeline parallelism is enabled (degree > 1)."""

        return self.pipeline_parallel > 1

    @property
    def has_data_parallel_replicate(self) -> bool:
        """Checks if data parallel replication is enabled (degree > 1)."""

        return self.data_parallel_replicate > 1

    @property
    def has_data_parallel_shard(self) -> bool:
        """Checks if data parallel sharding is enabled (degree > 1)."""

        return self.data_parallel_shard > 1

    @property
    def has_context_parallel_replicate(self) -> bool:
        return self.context_parallel_replicate > 1

    @property
    def has_context_parallel_shard(self) -> bool:
        return self.context_parallel_shard > 1

    @property
    def has_tensor_parallel(self) -> bool:
        return self.tensor_parallel > 1

    @property
    def has_expert_parallel(self) -> bool:
        """Checks if expert parallelism is enabled (degree > 1)."""
        return self.expert_parallel > 1

    @property
    def is_distributed(self) -> bool:
        """Checks if any form of parallelism is enabled."""

        return (
                self.has_pipeline_parallel or
                self.has_data_parallel_replicate or
                self.has_data_parallel_shard or
                self.has_context_parallel_shard or
                self.has_context_parallel_replicate or
                self.has_expert_parallel or
                self.has_tensor_parallel
        )

    @model_validator(mode="after")
    def _check_ep_divisibility(self) -> Self:
        """Validates that DP/CP/TP dimensions can support the requested EP/ETP degrees."""
        dp_cp_tp_degree = (
            self.data_parallel_shard *
            self.data_parallel_replicate *
            self.context_parallel_shard *
            self.context_parallel_replicate *
            self.tensor_parallel
        )
        ep_degree = self.expert_parallel

        if dp_cp_tp_degree % ep_degree != 0:
            raise ValueError(
                f"Total data/context/tensor parallelism degree ({dp_cp_tp_degree}) must be divisible by "
                f"total expert parallelism degree ({ep_degree})."
            )
        return self

    def build(self, log_level: int = logging.INFO) -> "DistributedContext":
        """
        Initializes the DistributedContext using these parameters.

        Returns:
            A new DistributedContext instance containing the initialized device meshes.
        """

        return DistributedContext(self, log_level)

class DistributedContext:
    """
    Acts as the single source of truth for the distributed execution environment.

    It acts as the central repository for the distributed configuration, managing the creation
    and synchronization of PyTorch DeviceMeshes for different domains (Regular domain, Expert Parallel domain, ...).

    All assertions regarding rank placement, group memberships, and parallel topology
    must be derived from this context to ensure consistency.
    """

    def __init__(self, params: "DeviceMeshParameters", log_level: int):
        self._params = params

        if params.is_distributed:
            meshes = _build_mesh_domains(params)
            regular_mesh = meshes[REGULAR_DOMAIN]

            self._meshes = meshes
            self._num_nodes = regular_mesh.size() // torch.cuda.device_count()
            self._logger = build_dist_logger(
                f'pp:{regular_mesh.get_local_rank("pp")}-'
                f'dpr:{regular_mesh.get_local_rank("dp_replicate")}-'
                f'dps:{regular_mesh.get_local_rank("dp_shard")}-'
                f'cps:{regular_mesh.get_local_rank("cp_shard")}-'
                f'cpr:{regular_mesh.get_local_rank("cp_replicate")}-'
                f'tp:{regular_mesh.get_local_rank("tp")}',
                level=log_level
            )
        else:
            self._meshes = {}
            self._num_nodes = 1
            self._logger = build_dist_logger("local", level=log_level)

        self._local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        self._global_rank = int(os.environ.get("RANK", "0"))

        self._node_rank = self._global_rank // torch.cuda.device_count()

        self._master_addr = _resolve_master_addr()
        self._current_device = torch.device("cuda")

        torch.cuda.set_device(self._local_rank)

    @property
    def logger(self) -> logging.Logger:
        """Returns the logger instance configured for distributed logging."""

        return self._logger

    def mesh_for(self, domain: str) -> DeviceMesh:
        """
        Returns the device mesh view associated with a specific logical domain.

        Available Domains and Dimensions:
            *   `regular` (`REGULAR_DOMAIN`): The most granular mesh for fully decomposed parallelism.
                Dimensions: ``('pp', 'dp_replicate', 'dp_shard', 'cp_shard', 'cp_replicate', 'tp')``
            *   `expert` (`EXPERT_DOMAIN`): Mesh optimized for distributing MoE (Mixture of Experts) layers.
                Dimensions: ``('pp', 'replicate', 'ep')``
            *   `dense` (`DENSE_DOMAIN`): Mesh optimized for distributing dense layers.
                Dimensions: ``('pp', 'dp_replicate', 'dp_cp_shard', 'cp_replicate', 'tp')``
            *   `batch` (`BATCH_DOMAIN`): Mesh optimized for distributing input data.
                Dimensions: ``('pp', 'dp', 'cp', 'tp')``
            *   `flat` (`FLAT_DOMAIN`): Mesh containing a single dimension with all the processes.
                Dimensions: ``('world')``

        Args:
            domain: The name of the domain to retrieve.

        Returns:
            The PyTorch DeviceMesh configured for the requested domain.

        Raises:
            ValueError: If the specified domain does not exist.
        """

        if domain not in self._meshes:
            raise ValueError(f"Domain {domain} does not exist")
        return self._meshes[domain]

    @property
    def is_main_process(self) -> bool:
        """Checks if the current process is the global rank 0."""

        return self._global_rank == 0

    @property
    def is_local_main_process(self) -> bool:
        """Checks if the current process is the rank 0 on the specific node."""

        return self._local_rank == 0

    def wait_world(self):
        """Blocks process execution until all ranks reach this point."""

        torch.distributed.barrier(device_ids=[torch.cuda.current_device()])
        torch.cuda.synchronize()

    def set_timeout(self, timeout_seconds: float):
        """
        Updates the NCCL/process group timeout for all underlying meshes.

        Args:
            timeout_seconds: New timeout duration in seconds.
        """

        self.logger.info(f"Setting global timeout to {timeout_seconds} seconds")
        self.wait_world()

        groups: list[torch.distributed.ProcessGroup | None] = [None]
        for mesh in self._meshes.values():
            for dim in range(mesh.ndim):
                groups.append(mesh.get_group(dim))

        for group in groups:
            torch.distributed.distributed_c10d._set_pg_timeout(datetime.timedelta(seconds=timeout_seconds), group)  # noqa: SLF001

    @contextmanager
    def local_main_process_first(self):
        """
        Context manager that executes the block on the local main process first.

        Other local ranks wait at the entrance. The local main process waits at the
        exit to synchronize before continuing.
        """
        if not self.is_local_main_process:
            self.wait_world()

        yield

        if self.is_local_main_process:
            self.wait_world()

    @contextmanager
    def main_process_first(self):
        """
        Context manager that executes the block on the global main process first.

        All other ranks wait at the entrance. The global main process waits at the
        exit to synchronize before continuing.
        """

        if not self.is_main_process:
            self.wait_world()

        yield

        if self.is_main_process:
            self.wait_world()

    @property
    def current_device(self) -> torch.device:
        """Returns the CUDA device associated with this rank."""

        return self._current_device

    @property
    def mesh_params(self) -> "DeviceMeshParameters":
        """Returns the parameters used to initialize this context."""

        return self._params

    @property
    def master_addr(self) -> str:
        """Returns the IP address or domain name of the master node."""

        return self._master_addr

    @property
    def node_rank(self) -> int:
        """Returns the index of the node this process is running on."""

        return self._node_rank

    @property
    def num_nodes(self) -> int:
        """Returns the total number of nodes in the cluster."""

        return self._num_nodes