class PipelineStage:
    """
    Represents a single structural stage in a Pipelined Model.

    This class acts as an orchestrator that combines `StageCommunicationHandler` (for I/O)
    and `Forward/BackwardComputeHandler` (for execution). It abstracts away the complexity
    of buffer management, distributed communication, and gradient calculation from the scheduler.
    """

    def __init__(
            self,
            info: PipelineStageInfo,
            module: nn.Module,
            group: dist.ProcessGroup,
            stage_to_host_topology: dict[int, int]
    ):
        """
        Constructs a PipelineStage object.

        Args:
            info: Metadata about the stage (index, total stages).
            module: The PyTorch module executed by this stage.
            group: The distributed process group for pipeline communications.
            stage_to_host_topology: Dict mapping stage ID to PP rank hosting it.
        """

        self._info = info
        self._module = module
        self._group = group
        self._stage_to_host_topology = stage_to_host_topology

        self._has_backward = False

        self._forward_comm: StageCommunicationHandler | None = None
        self._backward_comm: StageCommunicationHandler | None = None

        self._forward_comp = ForwardComputeHandler(
            stage_index=info.current_stage,
            module=module
        )
        self._backward_comp = BackwardComputeHandler(
            stage_index=info.current_stage,
            module=module
        )

    @property
    def info(self) -> PipelineStageInfo:
        return self._info

    def configure_buffers(
            self,
            num_microbatches: int,
            has_backward: bool,
            pipeline_inputs: dict[str, torch.Tensor]
    ):
        """
        Initializes the communication handlers and buffers for the stage.

        This must be called before execution to establish P2P buffer sizes and directions.

        Args:
            num_microbatches: Total number of microbatches to process.
            has_backward: Does this pipeline stage should store info for a backward pass
            pipeline_inputs: Pipeline input data.
        """

        self._has_backward = has_backward

        prev_stage_idx = None if self._info.is_current_stage_first else self._info.current_stage - 1
        next_stage_idx = None if self._info.is_current_stage_last else self._info.current_stage + 1

        with torch.device("meta"):
            if not isinstance(self._module, ModuleSupportsPipelining):
                raise TypeError("Module does not implement ModuleSupportsPipelining protocol")
            inputs_meta = self._module.infer_stage_inputs_from_pipeline_inputs(
                inputs=pipeline_inputs,
                n_microbatches=num_microbatches
            )
            outputs_meta = self._module.infer_stage_outputs_from_pipeline_inputs(
                inputs=pipeline_inputs,
                n_microbatches=num_microbatches
            )

        self._forward_comm = StageCommunicationHandler(
            name="fwd",
            stage_index=self._info.current_stage,
            num_microbatches=num_microbatches,
            input_stage_index=prev_stage_idx,
            input_args=inputs_meta,
            output_stage_index=next_stage_idx,
            output_args=outputs_meta,
            group=self._group,
            stage_idx_to_host_rank=self._stage_to_host_topology
        )
        self._forward_comm.set_input_requires_grad_(requires_grad=has_backward)

        if has_backward:
            # for grad - current stage receives OUTPUTS as inputs and sends INPUTS as outputs
            # because it is reversed forward
            self._backward_comm = StageCommunicationHandler(
                name="bwd",
                stage_index=self._info.current_stage,
                num_microbatches=num_microbatches,
                input_stage_index=next_stage_idx,
                input_args=outputs_meta,
                output_stage_index=prev_stage_idx,
                output_args=inputs_meta,
                group=self._group,
                stage_idx_to_host_rank=self._stage_to_host_topology
            )
        else:
            self._backward_comm = None

    def set_local_fwd_input(self, inputs: dict[str, torch.Tensor], microbatch_index: int):
        """
        Sets local forward inputs manually.

        Used for the V-shape schedulers.
        """

        if self._forward_comm is None:
            raise ValueError("You must configure stage buffers first")

        self._forward_comm.set_inputs_local(inputs, microbatch_index)

    def get_local_fwd_output(self, microbatch_index: int) -> dict[str, torch.Tensor]:
        return self._forward_comp.get_outputs(microbatch_index)

    def pop_local_bwd_output(self, microbatch_index: int) -> dict[str, torch.Tensor]:
        """
        Retrieves local backward outputs (gradients).
        """

        if not self._has_backward:
            raise ValueError()

        return self._backward_comp.pop_for_sending(microbatch_index)

    def set_local_bwd_input(self, inputs: dict[str, torch.Tensor], microbatch_index: int):
        """
        Sets local backward inputs (output gradients) manually.
        """

        if not self._has_backward:
            raise ValueError()

        if self._backward_comm is None:
            raise ValueError("You must configure stage buffers first")

        self._backward_comm.set_inputs_local(inputs, microbatch_index)

    def get_fwd_recv_ops(self, microbatch_index: int) -> list[dist.P2POp]:
        """Returns P2P ops to receive forward inputs for the given microbatch."""

        if self._forward_comm is None:
            raise ValueError("You must configure stage buffers first")

        return self._forward_comm.create_receive_ops(microbatch_index)

    def get_fwd_send_ops(self, microbatch_index: int) -> list[dist.P2POp]:
        """Returns P2P ops to send forward outputs for the given microbatch."""

        if self._forward_comm is None:
            raise ValueError("You must configure stage buffers first")

        fwd_result = self._forward_comp.get_outputs(microbatch_index)
        return self._forward_comm.create_send_ops(fwd_result)

    def get_bwd_recv_ops(self, microbatch_index: int) -> list[dist.P2POp]:
        """Returns P2P ops to receive backward gradients for the given microbatch."""

        if not self._has_backward:
            return []

        if self._backward_comm is None:
            raise ValueError("You must configure stage buffers first")

        return self._backward_comm.create_receive_ops(microbatch_index)

    def get_bwd_send_ops(self, microbatch_index: int) -> list[dist.P2POp]:
        """Returns P2P ops to send backward gradients for the given microbatch."""

        if not self._has_backward:
            return []

        if self._backward_comm is None:
            raise ValueError("You must configure stage buffers first")

        bwd_result = self._backward_comp.pop_for_sending(microbatch_index)
        return self._backward_comm.create_send_ops(bwd_result)

    def forward_one_chunk(
            self,
            microbatch_index: int,
            pipeline_inputs: dict[str, torch.Tensor],
            pipeline_kwargs: dict[str, Any] | None = None,
    ):
        """
        Executes a forward pass for a single microbatch chunk.

        Fetches inputs from the communication buffer (or `pipeline_inputs` if first stage),
        runs the computation, and caches the result.

        Args:
            microbatch_index: The microbatch index.
            pipeline_inputs: Inputs provided locally (only used if this is the first stage).
            pipeline_kwargs: Additional arguments for the module.

        Returns:
            The output tensors of the forward pass.
        """

        if self._forward_comm is None:
            raise ValueError("You must configure stage buffers first")

        if self._info.is_current_stage_first:
            inputs = pipeline_inputs
        else:
            inputs = self._forward_comm.get_inputs(microbatch_index)

        kwargs = pipeline_kwargs or {}

        self._forward_comp.run(
            microbatch_index=microbatch_index,
            inputs=inputs,
            kwargs=kwargs
        )

    def backward_one_chunk(
            self,
            microbatch_index: int,
            loss: torch.Tensor | None = None,
            full_backward: bool = True
    ):
        """
        Executes a backward pass for a single microbatch chunk.

        Can perform either a full backward or just the input gradients (if `full_backward=False`).
        It fetches required data from forward cache and communication buffers.

        Args:
            microbatch_index: The microbatch index.
            loss: The loss tensor (only used if this is the last stage).
            full_backward: If True, computes grads for inputs and weights. If False, only for inputs.
        """

        if not self._has_backward:
            raise ValueError()

        if self._backward_comm is None:
            raise ValueError("You must configure stage buffers first")

        inputs, fwd_outputs = self._forward_comp.pop_inputs_outputs(microbatch_index)

        outputs: dict[str, torch.Tensor]
        outputs_grad: dict[str, torch.Tensor] | None

        if self._info.is_current_stage_last:
            if loss is None:
                raise ValueError("Cannot perform backward on last stage without loss specified")
            outputs = {"loss": loss}
            outputs_grad = None
        else:
            outputs = fwd_outputs
            outputs_grad = self._backward_comm.get_inputs(microbatch_index)

        if full_backward:
            self._backward_comp.backward_full(
                microbatch_index=microbatch_index,
                inputs=inputs,
                outputs=outputs,
                outputs_grad=outputs_grad
            )
        else:
            self._backward_comp.backward_input(
                microbatch_index=microbatch_index,
                inputs=inputs,
                outputs=outputs,
                outputs_grad=outputs_grad
            )

        if self._info.is_current_stage_last and not self._info.is_current_stage_first:
            for t in fwd_outputs.values():
                if not t._is_view():  # noqa: SLF001
                    t.detach_()

    def backward_weight_one_chunk(self, microbatch_index: int):
        """
        Executes the weight gradient accumulation part of the backward pass.

        This assumes `backward_one_chunk(..., full_backward=False)` was already called
        for this microbatch.

        Args:
            microbatch_index: The microbatch index.
        """

        if not self._has_backward:
            raise ValueError()

        self._backward_comp.backward_weight(microbatch_index=microbatch_index)

    def reset(self):
        """Resets the internal state of communication handlers, clearing gradients on buffers."""

        if self._forward_comm is not None:
            self._forward_comm.reset()
        if self._backward_comm is not None:
            self._backward_comm.reset()

class ActionBase(abc.ABC):
    """
    Abstract base class for all pipeline schedule actions.

    An action represents an atomic unit of work in a pipeline schedule,
    such as computing a microbatch or sending/receiving a tensor.
    """

    @abc.abstractmethod
    def apply(self, ctx: ActionContext):
        """
        Executes the action logic using the provided context.

        Args:
            ctx: The runtime context containing stages, data, and communication handlers.
        """

        ...

    @property
    @abc.abstractmethod
    def work_type(self) -> ActionWorkType:
        """Returns the classification of work this action performs."""
        ...

    @property
    @abc.abstractmethod
    def has_backward_work(self) -> bool:
        """Returns True if this action involves backward pass computations."""
        ...

    @abc.abstractmethod
    def __str__(self) -> str:
        """Returns a short string representation of the action for logging/visualization."""
        ...

@dataclasses.dataclass(frozen=True, slots=True)
class BackwardFullInputComputeAction(ActionBase):
    """
    Action to perform backward computation with respect to inputs.

    Attributes:
        stage_idx: The integer index of the pipeline stage.
        microbatch_idx: The integer index of the microbatch to compute.
        full_backward: If True, performs a full backward pass including inputs
            and weights. If False, may only compute gradients w.r.t inputs
            (depending on schedule implementation).
    """

    stage_idx: int
    microbatch_idx: int
    full_backward: bool

    def apply(self, ctx: ActionContext):
        # todo unshard
        stage = ctx.stages[self.stage_idx]

        if not stage.info.is_current_stage_last and self.stage_idx + 1 not in ctx.stages:
            ctx.communications.wait_bwd_recv(self.stage_idx, self.microbatch_idx)

        if stage.info.is_current_stage_last and ctx.loss is not None:
            loss = ctx.loss.acquire_loss(self.microbatch_idx)
        else:
            loss = None

        stage.backward_one_chunk(
            microbatch_index=self.microbatch_idx,
            full_backward=self.full_backward,
            loss=loss
        )

        if not stage.info.is_current_stage_first and self.stage_idx - 1 in ctx.stages:
            ctx.stages[self.stage_idx - 1].set_local_bwd_input(
                microbatch_index=self.microbatch_idx,
                inputs=stage.pop_local_bwd_output(self.microbatch_idx)
            )

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.compute

    @property
    def has_backward_work(self) -> bool:
        return True

    def __str__(self) -> str:
        letter = "B" if self.full_backward else "I"
        return f"{self.stage_idx}{letter}{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class BackwardReceiveAction(ActionBase):
    """
    Action to schedule a backward pass gradient receive operation.

    Attributes:
        stage_idx: The integer index of the pipeline stage expecting the receive operation.
        microbatch_idx: The integer index of the microbatch being received.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        ctx.communications.schedule_bwd_recv(self.stage_idx, self.microbatch_idx)

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.communicate

    @property
    def has_backward_work(self) -> bool:
        return True

    def __str__(self) -> str:
        return f"{self.stage_idx}RECV_B{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class BackwardSendAction(ActionBase):
    """
    Action to schedule a backward pass gradient send operation.

    Attributes:
        stage_idx: The integer index of the pipeline stage initiating the send operation.
        microbatch_idx: The integer index of the microbatch being sent.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        ctx.communications.schedule_bwd_send(self.stage_idx, self.microbatch_idx)

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.communicate

    @property
    def has_backward_work(self) -> bool:
        return True

    def __str__(self) -> str:
        return f"{self.stage_idx}SEND_B{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class BackwardWeightComputeAction(ActionBase):
    """
    Action to perform gradient accumulation on weights.

    Attributes:
        stage_idx: The integer index of the pipeline stage.
        microbatch_idx: The integer index of the microbatch to compute.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        # todo unshard
        stage = ctx.stages[self.stage_idx]

        stage.backward_weight_one_chunk(
            microbatch_index=self.microbatch_idx
        )

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.compute

    @property
    def has_backward_work(self) -> bool:
        return True

    def __str__(self) -> str:
        return f"{self.stage_idx}W{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class ComposeAction(ActionBase):
    """
    Composite action scheduling multiple sub-actions sequentially.

    Used for forward/backward overlapping.

    Attributes:
        actions: A tuple of sub-actions to be executed sequentially.
    """

    actions: tuple[ActionBase, ...]

    def apply(self, ctx: ActionContext):
        for act in self.actions:
            act.apply(ctx)

    @property
    def work_type(self) -> ActionWorkType:
        sub_work_types = {x.work_type for x in self.actions}
        if len(sub_work_types) != 1:
            raise ValueError("")
        return next(iter(sub_work_types))

    @property
    def has_backward_work(self) -> bool:
        return any(x.has_backward_work for x in self.actions)

    def __str__(self) -> str:
        return "|".join(map(str, self.actions))

@dataclasses.dataclass(frozen=True, slots=True)
class ForwardComputeAction(ActionBase):
    """
    Action to perform forward computation for a specific microbatch.

    Attributes:
        stage_idx: The integer index of the pipeline stage.
        microbatch_idx: The integer index of the microbatch to compute.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        # todo check unsharded
        stage = ctx.stages[self.stage_idx]

        if not stage.info.is_current_stage_first and self.stage_idx - 1 not in ctx.stages:
            ctx.communications.wait_fwd_recv(self.stage_idx, self.microbatch_idx)

        stage.forward_one_chunk(
            microbatch_index=self.microbatch_idx,
            pipeline_inputs=ctx.pipeline_inputs_microbatches[self.microbatch_idx],
            pipeline_kwargs=ctx.pipeline_kwargs_microbatches[self.microbatch_idx]
        )
        result = stage.get_local_fwd_output(self.microbatch_idx)

        if stage.info.is_current_stage_last and ctx.loss is not None:
            ctx.loss.compute_loss(result, self.microbatch_idx)

        if not stage.info.is_current_stage_last and self.stage_idx + 1 in ctx.stages:
            ctx.stages[self.stage_idx + 1].set_local_fwd_input(
                inputs=result,
                microbatch_index=self.microbatch_idx
            )

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.compute

    @property
    def has_backward_work(self) -> bool:
        return False

    def __str__(self) -> str:
        return f"{self.stage_idx}F{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class ForwardReceiveAction(ActionBase):
    """
    Action to schedule a forward pass tensor receive operation.

    Attributes:
        stage_idx: The integer index of the pipeline stage expecting the receive operation.
        microbatch_idx: The integer index of the microbatch being received.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        ctx.communications.schedule_fwd_recv(self.stage_idx, self.microbatch_idx)

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.communicate

    @property
    def has_backward_work(self) -> bool:
        return True

    def __str__(self) -> str:
        return f"{self.stage_idx}RECV_F{self.microbatch_idx}"

@dataclasses.dataclass(frozen=True, slots=True)
class ForwardSendAction(ActionBase):
    """
    Action to schedule a forward pass tensor send operation.

    Attributes:
        stage_idx: The integer index of the pipeline stage initiating the send operation.
        microbatch_idx: The integer index of the microbatch being sent.
    """

    stage_idx: int
    microbatch_idx: int

    def apply(self, ctx: ActionContext):
        ctx.communications.schedule_fwd_send(self.stage_idx, self.microbatch_idx)

    @property
    def work_type(self) -> ActionWorkType:
        return ActionWorkType.communicate

    @property
    def has_backward_work(self) -> bool:
        return False

    def __str__(self) -> str:
        return f"{self.stage_idx}SEND_F{self.microbatch_idx}"

class PipelineScheduleExecutor(PipelineSchedule):
    """Executes a defined pipeline schedule by interpreting a sequence of actions."""

    def __init__(
            self,
            dist_context: DistributedContext,
            stages: list[PipelineStage],
            num_microbatches: int,
            loss_fn: LossFn | None,
            program: dict[int, list[ActionBase]],
            sharding_spec: PipelineShardingSpec
    ):
        """
        Constructs the schedule executor.

        Args:
            dist_context: The distributed context.
            stages: List of stages managed by this executor.
            num_microbatches: Number of microbatches the global batch is split.
            loss_fn: Function to compute loss.
            program: The execution plan mapping rank ID to a list of actions.
            sharding_spec: Sharding specification for input and output tensors.
        """

        self._dist_ctx = dist_context
        self._stages = {stage.info.current_stage: stage for stage in stages}
        self._num_microbatches = num_microbatches
        self._program = program

        self._has_backward = any(any(
            action.has_backward_work for action in sub_program
        ) for sub_program in program.values())

        self._comm_handler = PipelineCommunicationHandler(self._stages)
        if loss_fn is None:
            self._loss_handler = None
        else:
            self._loss_handler = PipelineLossHandler(loss_fn)

        # these could be late-initialized on configure_buffers \/
        self._input_data_sharding_spec = sharding_spec.input_data
        self._input_kwargs_sharding_spec = sharding_spec.input_kwargs

    def configure_buffers(self, inputs: dict[str, torch.Tensor], kwargs: dict[str, Any]):
        if self._input_data_sharding_spec is None:
            self._input_data_sharding_spec = shard_spec_on_dim(inputs, dim=0)
        if self._input_kwargs_sharding_spec is None:
            self._input_kwargs_sharding_spec = shard_spec_on_dim(kwargs, dim=0)

        for stage in self._stages.values():
            stage.configure_buffers(
                num_microbatches=self._num_microbatches,
                pipeline_inputs=inputs,
                has_backward=self._has_backward
            )

    def step(self, inputs: dict[str, torch.Tensor], kwargs: dict[str, Any]):
        if self._input_data_sharding_spec is None or self._input_kwargs_sharding_spec is None:
            raise ValueError("Please configure sharding specs first")

        self._dist_ctx.logger.debug("Begin pipeline step")
        pp_group = self._dist_ctx.mesh_for(REGULAR_DOMAIN).get_group("pp")

        for stage in self._stages.values():
            stage.reset()

        # Shard inputs and kwargs to microbatches
        inputs_shard = shard_tree(
            inputs,
            num_shards=self._num_microbatches,
            sharding_spec=self._input_data_sharding_spec,
            enforce_even_split=True
        )
        kwargs_shard = shard_tree(
            kwargs,
            num_shards=self._num_microbatches,
            sharding_spec=self._input_kwargs_sharding_spec,
            enforce_even_split=True
        )

        my_program = self._program[pp_group.rank()]

        for action in my_program:
            with record_function(str(action)):
                self._dist_ctx.logger.debug(f"Running pipeline action {action}")
                action.apply(ActionContext(
                    loss=self._loss_handler,
                    stages=self._stages,
                    communications=self._comm_handler,
                    pipeline_inputs_microbatches=inputs_shard,
                    pipeline_kwargs_microbatches=kwargs_shard
                ))

        self._dist_ctx.logger.debug("Waiting for potentially hanging PP send comms")
        self._comm_handler.wait_send_all()  # finalize just in case
        self._dist_ctx.logger.debug("End pipeline step")

class PipelineProgramBuilder(abc.ABC):
    """Abstract interface for building pipeline execution schedules."""

    @abc.abstractmethod
    def compose(self, num_microbatches: int, pp_size: int) -> dict[int, list[ActionBase]]:
        """
        Generates the execution program for all ranks in the pipeline.

        Args:
            num_microbatches: Number of microbatches per step.
            pp_size: Number of pipeline parallel ranks.

        Returns:
            A dictionary mapping rank indices to their list of sequential actions.
        """
        ...

    @property
    @abc.abstractmethod
    def num_stages_per_rank(self) -> int:
        """Returns the number of model stages designated for each rank."""

        ...

    @property
    @abc.abstractmethod
    def topology_style(self) -> ScheduleStyle:
        """Returns the topology style strategy used to assign stages to ranks."""
        ...

class ScheduleStyle(StrEnum):
    """
    Defines the strategy for mapping logical stages to physical ranks.

    Attributes:
        loop: Assigns stages in a round-robin circular fashion (mod pp_size).
        v: Assigns stages in a zig-zag V-shape pattern. Useful for interleaved 1F1B schedules.
    """

    loop = "loop"
    v = "v"