"""Packed-input H3 transformer core, independent of ComfyUI node execution."""

import torch
from torch import nn
from typing import TYPE_CHECKING

from .attention import DEFAULT_ATTENTION_BACKEND
from .backbone import H3DenoiserBackbone
from .checkpoint import H3Checkpoint
from .final import H3FinalLayer

if TYPE_CHECKING:
    from .distributed import SequenceParallelContext


class H3PackedDenoiser(nn.Module):
    """Run the H3 transformer once its Ref2VA payload has been packed into hidden rows."""

    def __init__(self, backbone: H3DenoiserBackbone, final_layer: H3FinalLayer):
        super().__init__()
        self.backbone = backbone
        self.final_layer = final_layer

    @classmethod
    def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = DEFAULT_ATTENTION_BACKEND):
        return cls(
            H3DenoiserBackbone.from_checkpoint(checkpoint, output_dtype=output_dtype, attention_backend=attention_backend),
            H3FinalLayer.from_checkpoint(checkpoint, output_dtype=output_dtype),
        )

    def forward(
        self,
        hidden: torch.Tensor,
        timesteps: torch.Tensor,
        position_ids: torch.Tensor,
        segments: list[tuple[int, int, int]],
        video_segment: tuple[int, int, int],
        audio_segment: tuple[int, int, int],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        hidden = self.backbone(hidden, timesteps, position_ids, segments)
        return self.final_layer(hidden, timesteps, video_segment, audio_segment)

    def forward_sequence_parallel(
        self,
        full_hidden: torch.Tensor,
        timesteps: torch.Tensor,
        full_position_ids: torch.Tensor,
        segments: list[tuple[int, int, int]],
        video_segment: tuple[int, int, int],
        audio_segment: tuple[int, int, int],
        context: "SequenceParallelContext",
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run all 50 blocks with token-sharded activations and Ulysses attention."""
        if full_hidden.shape[0] != context.sequence_length:
            raise ValueError("packed hidden length does not match sequence-parallel context")
        start, stop = context.local_token_range
        local_hidden = full_hidden[start:stop].contiguous()
        local_positions = full_position_ids[start:stop].contiguous()
        del full_hidden, full_position_ids
        local_hidden = self.backbone(
            local_hidden,
            timesteps,
            local_positions,
            segments,
            sequence_parallel=context,
        )
        return self.final_layer.forward_sequence_parallel(
            local_hidden,
            timesteps,
            video_segment,
            audio_segment,
            context,
        )

    def forward_tensor_parallel(
        self,
        full_hidden: torch.Tensor,
        timesteps: torch.Tensor,
        full_position_ids: torch.Tensor,
        segments: list[tuple[int, int, int]],
        video_segment: tuple[int, int, int],
        audio_segment: tuple[int, int, int],
        context: "SequenceParallelContext",
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run TP-sharded NVFP4 linears with ragged sequence-sharded residuals."""
        if full_hidden.shape[0] != context.sequence_length:
            raise ValueError("packed hidden length does not match tensor-parallel context")
        start, stop = context.local_token_range
        local_hidden = full_hidden[start:stop].contiguous()
        local_positions = full_position_ids[start:stop].contiguous()
        del full_hidden, full_position_ids
        local_hidden = self.backbone(
            local_hidden,
            timesteps,
            local_positions,
            segments,
            tensor_parallel=context,
        )
        return self.final_layer.forward_sequence_parallel(
            local_hidden,
            timesteps,
            video_segment,
            audio_segment,
            context,
        )