# 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 inspect
from functools import wraps
from typing import TYPE_CHECKING

from ..modeling_outputs import CausalLMOutputWithPast
from ..utils import is_torch_available, is_torch_distributed_available


if TYPE_CHECKING:
    import torch.nn as nn

if is_torch_available():
    import torch
    import torch.nn as nn

if is_torch_distributed_available():
    import torch.distributed as dist


def _bind_forward_kwargs(forward_signature: inspect.Signature, args: tuple, kwargs: dict) -> dict:
    bound = forward_signature.bind_partial(*args, **kwargs)
    bound.apply_defaults()
    return dict(bound.arguments)


class PipelineIdentityLayer(nn.Identity):
    """An identity layer replacing modules not owned by the current pipeline stage."""

    def __init__(self, *args, **kwargs):
        super().__init__()

    def forward(self, *args, **kwargs):
        """Return the first arg from args or the first value from kwargs as nothing is used."""
        return args[0] if args else next(iter(kwargs.values()))


class PipelineStage:
    """Pipeline-parallel stage metadata derived from a 1-D PP device mesh."""

    def __init__(self, pp_mesh: torch.distributed.device_mesh.DeviceMesh):
        self.pp_rank = pp_mesh.get_local_rank()
        self.pp_size = pp_mesh.size()
        self.pp_group = pp_mesh.get_group()
        self.pp_is_first_stage = self.pp_rank == 0
        self.pp_is_last_stage = self.pp_rank == self.pp_size - 1
        self.pp_prev_rank = self.pp_rank - 1 if self.pp_rank > 0 else None
        self.pp_next_rank = self.pp_rank + 1 if self.pp_rank < self.pp_size - 1 else None
        self.comm_on_cpu = dist.get_backend(self.pp_group) == "gloo"

    def communicate(
        self,
        operation: str,
        *,
        device: torch.device,
        dtype: torch.dtype,
        tensor: torch.Tensor | None = None,
        shape: tuple[int, ...] | None = None,
    ) -> torch.Tensor | None:
        """Point-to-point pipeline communication between adjacent stages."""
        comm_device = torch.device("cpu") if self.comm_on_cpu else device
        src = dest = None

        if operation == "recv_forward":
            if self.pp_is_first_stage:
                return None
            # Receive hidden states from the previous stage.
            src = self.pp_prev_rank
            # Shape is provided by the caller (derived from input_ids / inputs_embeds).
            tensor = torch.empty(shape, dtype=dtype, device=comm_device)

        elif operation == "send_forward":
            if self.pp_is_last_stage:
                return None
            # Send hidden states to the next stage.
            dest = self.pp_next_rank
            tensor = tensor.to(device=comm_device, dtype=dtype).contiguous()

        else:
            raise ValueError(f"Unsupported pipeline communication operation: {operation}")

        # Shared P2P: one isend/irecv with the adjacent rank.
        is_send = operation.startswith("send")
        peer_rank = dest if is_send else src
        op = dist.P2POp(dist.isend if is_send else dist.irecv, tensor, peer_rank, group=self.pp_group)
        # Wait for the communication to complete.
        for req in dist.batch_isend_irecv([op]):
            req.wait()
        if comm_device.type == "cuda":
            torch.cuda.synchronize()

        return None if is_send else tensor.to(device=device, dtype=dtype)

    def layer_range_for_rank(self, rank: int, num_layers: int) -> tuple[int, int]:
        """
        Return [start, end) owned by rank.

        Example with pp_size=3 and num_layers=32:
            layers_per_rank = 32 // 3 = 10
            rank 0 -> [0, 10)   # layers 0-9
            rank 1 -> [10, 20)  # layers 10-19
            rank 2 -> [20, 32)  # layers 20-31 (10 + 2 remainder)
        """
        # TODO(3outeille): Balance stages by layer type or parameter bytes instead of splitting solely by layer count.
        layers_per_rank = num_layers // self.pp_size
        start_layer = rank * layers_per_rank
        # last rank will always have more layers if uneven split. Else, it will have the same number of layers as the other ranks.
        end_layer = num_layers if rank == self.pp_size - 1 else start_layer + layers_per_rank
        return start_layer, end_layer

    def find_rank_for_key(self, key: str, num_layers: int, base_model_prefix: str) -> int | None:
        """Return the PP rank that owns a checkpoint parameter key, or ``None`` if unknown."""
        base_prefix = f"{base_model_prefix}."

        if key.startswith(f"{base_prefix}embed_tokens."):
            return 0

        if key.startswith(f"{base_prefix}norm.") or key.startswith("lm_head."):
            return self.pp_size - 1

        layers_prefix = f"{base_prefix}layers."
        if not key.startswith(layers_prefix):
            return None

        layer_idx = int(key.split(".")[2])
        for rank in range(self.pp_size):
            start_layer, end_layer = self.layer_range_for_rank(rank, num_layers)
            if start_layer <= layer_idx < end_layer:
                return rank

        return None

    def broadcast_from_last(
        self,
        tensor: torch.Tensor | None,
        *,
        dtype: torch.dtype,
        device: torch.device,
    ) -> torch.Tensor:
        """Broadcast logits from the last PP rank to every rank. Only the last stage computes logits, but every rank must return them so `generate()` works without PP-specific changes."""
        if self.pp_size <= 1:
            return tensor

        last_rank = self.pp_size - 1
        comm_device = torch.device("cpu") if self.comm_on_cpu else device
        # Logits are always (batch, seq_len, vocab_size).
        logits_ndim = 3

        # last rank sends the logits vector shape; other ranks receive it to size their buffer.
        if self.pp_is_last_stage:
            assert tensor.ndim == logits_ndim, (
                f"Expected logits with {logits_ndim} dims (batch, seq_len, vocab_size), got {tensor.ndim}"
            )
            logits = tensor.to(device=comm_device, dtype=dtype).contiguous()
            shape_msg = torch.tensor(list(logits.shape), dtype=torch.long, device=comm_device)
        else:
            shape_msg = torch.empty(logits_ndim, dtype=torch.long, device=comm_device)
        dist.broadcast(shape_msg, src=last_rank, group=self.pp_group)

        # last rank sends the actual logits value; other ranks receive into the pre-allocated buffer.
        if not self.pp_is_last_stage:
            logits = torch.empty(tuple(shape_msg.tolist()), dtype=dtype, device=comm_device)
        dist.broadcast(logits, src=last_rank, group=self.pp_group)

        return logits.to(device=device, dtype=dtype)


def apply_pipeline_parallelism(model: nn.Module, pp_mesh: torch.distributed.device_mesh.DeviceMesh) -> nn.Module:
    """Naive even split of `base_model.layers` across PP ranks."""
    # TODO(3outeille): involves pp_plan to do the split ?
    stage = PipelineStage(pp_mesh)
    model._pp_stage = stage

    base_model = getattr(model, model.base_model_prefix)
    layers = base_model.layers
    num_layers = len(layers)

    start_layer, end_layer = stage.layer_range_for_rank(stage.pp_rank, num_layers)
    tied = getattr(model.config, "tie_word_embeddings", False)

    # When tied, keep embed_tokens on the last stage too so _finalize_model_loading in modeling_utils.py can tie lm_head locally.
    keep_embed_tokens = stage.pp_is_first_stage or (tied and stage.pp_is_last_stage)
    if not keep_embed_tokens:
        base_model.embed_tokens = PipelineIdentityLayer()

    for layer_idx in range(num_layers):
        if layer_idx < start_layer or layer_idx >= end_layer:
            layers[layer_idx] = PipelineIdentityLayer()

    if not stage.pp_is_last_stage:
        base_model.norm = PipelineIdentityLayer()
        model.lm_head = PipelineIdentityLayer()

    # let _finalize_model_loading know that we want to tie the lm_head only in the last rank
    if tied and not stage.pp_is_last_stage:
        model.all_tied_weights_keys = {}

    # TODO(3outeille): dispatch to different pipeline parallelism schedules (gpipe, 1f1b, etc.)
    if not getattr(model, "_pp_forward_wrapped", False):
        original_forward = model.forward
        forward_signature = inspect.signature(original_forward)

        @wraps(original_forward)
        def pp_naive_forward(*args, **kwargs):
            return pipeline_parallel_naive_forward(model, original_forward, forward_signature, *args, **kwargs)

        # @wraps(original_forward)
        # def pp_gpipe_forward(*args, **kwargs):
        #     return pipeline_parallel_gpipe_forward(model, original_forward, *args, **kwargs)

        model.forward = pp_naive_forward
        model._pp_forward_wrapped = True

    return model


def _hidden_states_shape(fwd_kwargs: dict, hidden_size: int) -> tuple[int, ...]:
    if (inputs_embeds := fwd_kwargs.get("inputs_embeds")) is not None:
        return tuple(inputs_embeds.shape)
    if (input_ids := fwd_kwargs.get("input_ids")) is not None:
        return (*input_ids.shape[:2], hidden_size)
    raise ValueError("Cannot determine hidden states shape for pipeline recv_forward")


def _feed_hidden_states_as_input_embeds(fwd_kwargs: dict, hidden_states: torch.Tensor) -> dict:
    """Stage 0 uses input_ids; later stages use received hidden states as inputs_embeds."""
    fwd_kwargs.pop("input_ids", None)
    fwd_kwargs["inputs_embeds"] = hidden_states
    return fwd_kwargs


def pipeline_parallel_naive_forward(
    model: nn.Module, original_forward, forward_signature: inspect.Signature, *args, **kwargs
):
    stage = model._pp_stage

    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype

    # Bind the forward function to the arguments so that we can easily access them regardless of the number of arguments or the way we called the forward function.
    caller_kwargs = _bind_forward_kwargs(forward_signature, args, kwargs)
    fwd_kwargs = caller_kwargs

    # Non-first stages: recv activations from prev stage and use them as inputs_embeds.
    if not stage.pp_is_first_stage:
        shape = _hidden_states_shape(caller_kwargs, model.config.hidden_size)
        hidden_states = stage.communicate("recv_forward", device=device, dtype=dtype, shape=shape)
        fwd_kwargs = _feed_hidden_states_as_input_embeds(caller_kwargs, hidden_states)

    if stage.pp_is_last_stage:
        # Last stage: compute the logits.
        outputs = original_forward(**fwd_kwargs)
        logits, past_key_values = outputs.logits, outputs.past_key_values
    else:
        # Non-last stages: Compute and send the activations to the next stage.
        # We need the base models because hidden_states is getting slices in the causal_lm models + we want last_hidden_state
        base_model = getattr(model, model.base_model_prefix)
        base_kwargs = {k: v for k, v in fwd_kwargs.items() if k not in {"labels", "logits_to_keep"}}
        base_outputs = base_model(**base_kwargs)
        stage.communicate("send_forward", device=device, dtype=dtype, tensor=base_outputs.last_hidden_state)
        logits, past_key_values = None, base_outputs.past_key_values

    # Only the last stage computed the logits, so broadcast them to every rank (they will be waiting)
    # TODO(3outeille): very naive implementation to make it work seamless with generate().
    # In long term, we want only the last stage to hold the logits not every stages that will require changes in generate() directly
    logits = stage.broadcast_from_last(logits, dtype=dtype, device=device)

    return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)