# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import os
import warnings
from datetime import timedelta
from typing import TYPE_CHECKING, TypeGuard

from ..utils import is_torch_available, is_torch_distributed_available, is_torch_greater_or_equal, logging


logger = logging.get_logger(__name__)


if TYPE_CHECKING:
    from torch.distributed.device_mesh import DeviceMesh
    from torch.distributed.tensor import DTensor

    from .configuration_utils import DistributedConfig


if is_torch_available():
    import torch


def _is_torch_distributed_initialized() -> bool:
    if not is_torch_distributed_available():
        return False
    return torch.distributed.is_initialized()


def is_dtensor(obj: object) -> TypeGuard[DTensor]:
    if not is_torch_distributed_available():
        return False
    from torch.distributed.tensor import DTensor

    return isinstance(obj, DTensor)


def _get_torch_distributed_rank() -> int:
    if not _is_torch_distributed_initialized():
        return 0
    return torch.distributed.get_rank()


def _get_torch_distributed_world_size() -> int:
    if not _is_torch_distributed_initialized():
        return 1
    return torch.distributed.get_world_size()


def is_local_dist_rank_0() -> bool:
    return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0


DISTRIBUTED_BACKEND_MAP = {
    "cuda": "nccl",
    "cpu": "gloo",
    "xpu": "xccl",
    "hpu": "hccl",
    "neuron": "neuron",
    "tpu": "tpu_dist",
}


def get_distributed_backend(device_type: str) -> str | None:
    """Return the `torch.distributed` backend for `device_type`, or `None` if it has no dedicated one."""
    dist_backend = DISTRIBUTED_BACKEND_MAP.get(device_type)
    if dist_backend is None:
        raise ValueError(f"No distributed backend found for device type '{device_type}'")
    return dist_backend


def _bind_local_device(device_type: str, local_rank: int) -> torch.device:
    """Bind this process to its accelerator and return the device (no binding on CPU)."""
    if device_type == "cpu":
        return torch.device(device_type)
    # torch_tpu gives each process a single chip, always at index 0
    device_index = torch.accelerator.current_device_index() if device_type == "tpu" else local_rank
    torch.accelerator.set_device_index(device_index)
    return torch.device(device_type, device_index)


def _ensure_torch_distributed(device_type: str | None = None):
    """Initialize torch.distributed if not already initialized.

    If `device_type` is not given, it is detected from the current accelerator.
    """
    if not torch.distributed.is_initialized():
        if device_type is None:
            device_type = torch._C._get_accelerator().type
        if device_type == "mps":
            logger.warning_once(
                "PyTorch's built-in DeviceMesh/DTensor stack does not support an MPS mesh. Falling back to CPU."
            )
            device_type = "cpu"
        try:
            rank = int(os.environ["RANK"])
            local_rank = int(os.environ["LOCAL_RANK"])
            world_size = int(os.environ["WORLD_SIZE"])

            backend = get_distributed_backend(device_type)

            # Bind the accelerator before init so the process group is created with a
            # device_id, otherwise collectives like barrier() warn (and may spin up an
            # extra NCCL comm) about the missing device binding.
            device_id = None if device_type == "cpu" else _bind_local_device(device_type, local_rank)
            torch.distributed.init_process_group(
                backend=backend,
                rank=rank,
                world_size=world_size,
                device_id=device_id,
                # Sharded loading takes tens of minutes with high rank skew; the default 10-minute watchdog is too short
                timeout=timedelta(hours=2),
            )
        except Exception as e:
            raise OSError(
                "We tried to initialize torch.distributed for you, but it failed. Make "
                "sure you init torch distributed in your script to use distributed training."
            ) from e


def _distributed_barrier():
    """Barrier bound to the current accelerator device.

    Passing `device_ids` is required when the process group was initialized without a
    `device_id`; with it, the call is a no-op compared to plain `barrier()`. Safe to call
    when torch.distributed has not been initialized — returns immediately.
    """
    if not _is_torch_distributed_initialized():
        return
    device_type = torch._C._get_accelerator().type
    if device_type != "cpu":
        torch.distributed.barrier(device_ids=[torch.accelerator.current_device_index()])
    else:
        torch.distributed.barrier()


class TransformersDeviceMesh:
    """
    Holds the device meshes used by a model.

    dense layers and experts are sharded differently, so they need different views of the same ranks.

        dense  : (pp, fsdp, tp)     attention, dense MLPs, embeddings, lm_heads
        expert : (pp, efsdp, ep)    experts

    Both views cover the same world, so pp * fsdp * tp == pp * efsdp * ep.
    efsdp is not something you pick, it is whatever is left once ep is fixed. The relationship is as follow:
    efsdp = fsdp * tp / ep. It is the FSDP axis for expert weights same role `fsdp` plays for the dense params.
    There is no etp (expert tensor parallel) axis: with EP on, experts are never tensor-sharded. (for now)
    For experts, EP plays the role TP plays for dense layers (weights stay sharded during compute),
    and efsdp plays the role of fsdp (weights all-gathered before compute).

    If one were ever added, the relationship would become pp * efsdp * ep * etp == pp * fsdp * tp and efsdp would shrink by etp
    (efsdp = fsdp * tp / (ep * etp))

    When ep_size == tp_size, efsdp == fsdp (given the relationship efsdp = fsdp * tp / ep), thus same axis, same size and same groups of ranks.
    In that case experts could reuse the dense mesh's fsdp axis.
    When ep_size != tp_size, efsdp != fsdp, so we can't reuse the same mesh as they don't have the same groups of ranks
    This explains why experts need their own efsdp axis.

    Regarding ep value, one can decide to default the value to node width (8 on most machines) so that all-to-all never leaves the node.
    That has several implications on efsdp value given your setup:
    - On a single node, ep == fsdp * tp thus efsdp = 1, the axis does nothing.
    - On several nodes, we still keep ep at node width, since all-to-all across nodes is expensive.
    However, each node then holds a full copy of the expert group and efsdp is the number of copies, which is where FSDP happens for the experts
    i.e: 2 nodes x 8 GPUs -> efsdp = 16 / 8 = 2, one EP group per node, two copies, sharded over efsdp.
    """

    def __init__(self, dense_mesh: DeviceMesh, expert_mesh: DeviceMesh):
        self._dense_mesh = dense_mesh
        self._expert_mesh = expert_mesh

    def get_mesh(self, dims: str | tuple[str, ...]) -> DeviceMesh:
        """Select expert axes for `ep`/`efsdp`, otherwise dense axes; DeviceMesh handles slicing."""
        dims = (dims,) if isinstance(dims, str) else dims
        mesh = self._expert_mesh if "ep" in dims or "efsdp" in dims else self._dense_mesh
        return mesh[dims]


# Retained for the legacy transformers.integrations.tensor_parallel API.
def initialize_tensor_parallelism(
    tp_plan: str | dict[str, str] | None, tp_size: int | None = None, device_mesh=None, device_map=None
):
    r"""
    Sets up the device mesh and initialized the backend for tensor parallelism.
    This function is called when the model is loaded and the TP plan is set to 'auto'.
    """
    warnings.warn(
        "`initialize_tensor_parallelism` is deprecated and will be removed in a future release. "
        "Use `initialize_distributed_mesh` with a `DistributedConfig` instead.",
        FutureWarning,
        stacklevel=2,
    )
    if tp_size is not None and tp_plan is None:
        raise ValueError("tp_plan has to be set when tp_size is passed.")
    if tp_plan is not None and device_map is not None:
        raise ValueError("`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization.")
    if device_mesh is None:
        # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.
        device_type = torch._C._get_accelerator().type
        if device_type == "mps":
            logger.warning_once(
                "PyTorch's built-in DeviceMesh/DTensor stack does not support an MPS mesh. Falling back to CPU."
            )
            device_type = "cpu"
        if device_type != "cpu":
            tp_device = _bind_local_device(device_type, int(os.environ.get("LOCAL_RANK", 0)))
            device_map = tp_device
        else:
            tp_device = torch.device(device_type)
            device_map = device_type or {}

        device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))
    else:
        if device_mesh.ndim > 1:
            if "tp" not in device_mesh.mesh_dim_names:
                raise ValueError(
                    "When using `tp_plan` and n-d `device_mesh`, it must contain a 'tp' dimension. "
                    "Please provide a valid `device_mesh`."
                )
            device_mesh = device_mesh["tp"]
        device_map = _bind_local_device(device_mesh.device_type, int(os.environ.get("LOCAL_RANK", 0)))

    return device_map, device_mesh


def initialize_fully_sharded_data_parallelism(distributed_config: DistributedConfig):
    warnings.warn(
        "`initialize_fully_sharded_data_parallelism` is deprecated and will be removed in a future release. "
        "Use `initialize_distributed_mesh` with a `DistributedConfig` instead.",
        FutureWarning,
        stacklevel=2,
    )
    # `fully_shard` itself only needs torch>=2.6, but distributed checkpoint save/load
    # (DCP + HuggingFaceStorageWriter) needs 2.7, so that is the effective requirement.
    if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.7"):
        raise OSError("FSDP2 requires `torch>=2.7` (distributed checkpoint save/load).")

    device_type = torch._C._get_accelerator().type

    device_map = _bind_local_device(device_type, int(os.environ.get("LOCAL_RANK", 0)))

    fsdp_size = distributed_config.fsdp_size

    dims, names = [], []
    if fsdp_size > 1:
        dims.append(fsdp_size)
        names.append("fsdp")

    # Build the N-dimensional device mesh
    mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names))
    # If N > 1, create a flattened sub-mesh so all-reduces across the world mesh ae done in one collective
    if len(dims) > 1:
        mesh._flatten("_".join(names))

    return device_map, mesh


def initialize_distributed_mesh(
    distributed_config: DistributedConfig,
) -> tuple[torch.device | None, TransformersDeviceMesh | None]:
    """Create a device mesh containing every configured parallel dimension."""
    mesh_shape = (distributed_config.pp_size, distributed_config.fsdp_size, distributed_config.tp_size)
    if mesh_shape == (1, 1, 1):
        return None, None

    device_type = torch._C._get_accelerator().type
    if distributed_config.tp_size > 1 and device_type == "mps":
        raise RuntimeError("Tensor parallelism is not supported on MPS devices.")

    _ensure_torch_distributed(device_type)
    world_size = torch.distributed.get_world_size()
    expected_world_size = distributed_config.pp_size * distributed_config.fsdp_size * distributed_config.tp_size
    if expected_world_size != world_size:
        raise RuntimeError(
            f"The parallel mesh requires {expected_world_size} processes, but world_size is {world_size}."
        )

    device_map = _bind_local_device(device_type, int(os.environ.get("LOCAL_RANK", 0)))

    dense_mesh = torch.distributed.init_device_mesh(
        device_type,
        mesh_shape,
        mesh_dim_names=("pp", "fsdp", "tp"),
    )
    expert_mesh = torch.distributed.init_device_mesh(
        device_type,
        (distributed_config.pp_size, distributed_config.efsdp_size, distributed_config.ep_size),
        mesh_dim_names=("pp", "efsdp", "ep"),
    )
    return device_map, TransformersDeviceMesh(dense_mesh, expert_mesh)


def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
    """Gather FSDP-sharded params to full plain CPU tensors.

    Only rank 0 accumulates the result; other ranks return ``{}``.
    """
    if not is_torch_greater_or_equal("2.7"):
        raise OSError("Distributed checkpointing requires `torch>=2.7`.")

    # Import here because otherwise it emits a warning every time it's imported on some hardware - this keeps the warning from
    # being emitted if the function is not used
    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict

    options = StateDictOptions(full_state_dict=True, cpu_offload=True)
    full_state_dict = get_model_state_dict(model, options=options)
    if _get_torch_distributed_rank() == 0:
        return full_state_dict
    return {}


def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
    """Save model parameters as standard HF-format sharded safetensors using
    DCP + HuggingFaceStorageWriter with consolidation enabled.

    Every rank first writes its own shard in parallel under
    `<checkpoint_dir>/sharded/`, then a consolidation pass reads those shards
    and emits HF-compatible `model-*-of-N.safetensors` (+ index) at
    `<checkpoint_dir>/`. The result is a directory `from_pretrained` reads
    through its normal path — no special flag needed at load time.
    """
    if not is_torch_greater_or_equal("2.7"):
        raise OSError("Distributed checkpointing requires `torch>=2.7`.")

    # Import here because otherwise it emits a warning every time it's imported on some hardware - this keeps the warning from
    # being emitted if the function is not used
    import torch.distributed.checkpoint as dcp
    from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
    from torch.distributed.checkpoint.state_dict import get_model_state_dict

    state_dict = get_model_state_dict(model)
    dcp.save(
        state_dict,
        storage_writer=HuggingFaceStorageWriter(
            path=checkpoint_dir,
            save_distributed=True,
            enable_consolidation=True,
        ),
    )
    # Wait for rank 0 to finish writing the HF safetensors so other
    # ranks don't return (and hit `from_pretrained`) before the files exist.
    _distributed_barrier()


def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
    """Save optimizer state via DCP."""
    if not is_torch_greater_or_equal("2.7"):
        raise OSError("Distributed checkpointing requires `torch>=2.7`.")

    # Import here because otherwise it emits a warning every time it's imported on some hardware - this keeps the warning from
    # being emitted if the function is not used
    import torch.distributed.checkpoint as dcp
    from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict

    optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
    dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)


def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
    """Load optimizer state via DCP."""
    if not is_torch_greater_or_equal("2.7"):
        raise OSError("Distributed checkpointing requires `torch>=2.7`.")

    # Import here because otherwise it emits a warning every time it's imported on some hardware - this keeps the warning from
    # being emitted if the function is not used
    import torch.distributed.checkpoint as dcp
    from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict, set_optimizer_state_dict

    optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
    dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
    set_optimizer_state_dict(model, optimizer, optimizer_state_dict)