"""SCAIL-2 character-animation inference pipeline.

Two-stage distilled video generation that animates a character from a *driving*
video: the driving latent is concatenated into the token sequence with a RoPE
width offset (ΔW) and in-context mask channels route motion per character. This
mirrors the distilled pipeline (``distilled.py``) and lip-dub pipeline
(``lipdub.py``) structure, adding the SCAIL conditioning assembly.

Requirements:
- A SCAIL-trained checkpoint whose config declares ``mask_conditioning_channels``
  (so ``LTXModelConfigurator`` builds the widened ``patchify_proj`` automatically)
  plus the SCAIL LoRA adapter.
- A driving video (same target resolution) and a semantic mask tensor
  ``[K+1, F_pix, H_pix, W_pix]`` (channel 0 = environment switch, 1..K = binding
  slots), saved as a ``.pt`` file.

The SCAIL-specific conditioning assembly is factored into
``build_scail_conditionings`` so it can be unit-tested without the heavy models.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator

import torch

from ltx_core.components.noisers import GaussianNoiser
from ltx_core.conditioning import (
    ConditioningItem,
    DrivingMode,
    VideoConditionByDrivingLatent,
    VideoConditionByMaskChannels,
)
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import ImageConditioningInput, resolve_cli_params, scail_animation_arg_parser
from ltx_pipelines.utils.blocks import (
    DiffusionStage,
    ImageConditioner,
    PromptEncoder,
    VideoDecoder,
    VideoUpsampler,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS
from ltx_pipelines.utils.denoisers import SimpleDenoiser
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode


def build_scail_conditionings(
    driving_latent: torch.Tensor,
    masks: torch.Tensor | None,
    *,
    mode: DrivingMode = DrivingMode.ANIMATION,
    strength: float = 1.0,
    width_offset: float | None = None,
) -> list[ConditioningItem]:
    """Assemble the SCAIL driving + mask-channel conditioning items (order matters).

    The driving item is applied first (it appends the driving tokens at the tail); the mask item is
    applied second, writing its channels onto those trailing driving tokens (the noisy target keeps a
    zero mask). ``masks`` may be ``None`` to run driving-only (e.g. a model without mask channels).

    Args:
        driving_latent: Driving video latent ``[B, C, F, H, W]`` (same F/H/W as the target).
        masks: Semantic masks ``[B, K+1, F_pix, H_pix, W_pix]`` or ``None``.
        mode: SCAIL mode (animation/replacement).
        strength: Driving conditioning strength (1.0 keeps it clean/frozen).
        width_offset: ΔW in RoPE pixel-space width units (None = target pixel width).
    """
    conditionings: list[ConditioningItem] = [
        VideoConditionByDrivingLatent(
            latent=driving_latent,
            mode=mode,
            width_offset=width_offset,
            strength=strength,
        )
    ]
    if masks is not None:
        conditionings.append(VideoConditionByMaskChannels(masks=masks))
    return conditionings


def load_masks(mask_path: str, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Load a semantic mask tensor from a ``.pt`` file and shape it to ``[B, K+1, F_pix, H, W]``."""
    masks = torch.load(mask_path, map_location=device, weights_only=True)
    if isinstance(masks, dict):
        masks = masks["mask"]
    if masks.dim() == 4:  # [K+1, F, H, W] -> add batch
        masks = masks.unsqueeze(0)
    return masks.to(device=device, dtype=dtype)


class ScailAnimationPipeline:
    """Two-stage distilled SCAIL-2 character animation (driving video + in-context mask channels)."""

    def __init__(
        self,
        distilled_checkpoint_path: str,
        gemma_root: str,
        spatial_upsampler_path: str,
        scail_lora: LoraPathStrengthAndSDOps,
        device: torch.device | None = None,
        quantization: QuantizationPolicy | None = None,
        registry: Registry | None = None,
        compilation_config: CompilationConfig | None = None,
        offload_mode: OffloadMode = OffloadMode.NONE,
        alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
    ) -> None:
        self.device = device or get_device()
        self.dtype = torch.bfloat16

        self.prompt_encoder = PromptEncoder(
            distilled_checkpoint_path,
            gemma_root,
            self.dtype,
            self.device,
            registry=registry,
            offload_mode=offload_mode,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        self.image_conditioner = ImageConditioner(
            distilled_checkpoint_path,
            self.dtype,
            self.device,
            registry=registry,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        self.stage = DiffusionStage.from_checkpoint(
            distilled_checkpoint_path,
            self.dtype,
            self.device,
            loras=(scail_lora,),
            quantization=quantization,
            registry=registry,
            compilation_config=compilation_config,
            offload_mode=offload_mode,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        self.upsampler = VideoUpsampler(
            distilled_checkpoint_path,
            spatial_upsampler_path,
            self.dtype,
            self.device,
            registry=registry,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        self.video_decoder = VideoDecoder(
            distilled_checkpoint_path,
            self.dtype,
            self.device,
            registry=registry,
            alloc_trim_strategy=alloc_trim_strategy,
        )

    def _encode_driving_latent(
        self,
        driving_video_path: str,
        height: int,
        width: int,
        num_frames: int,
        video_encoder: VideoEncoder,
        tiling_config: TilingConfig | None,
    ) -> torch.Tensor:
        """Decode + encode the driving video into a VAE latent at the target resolution."""
        frame_gen = decode_video_by_frame(path=driving_video_path, frame_cap=num_frames, device=self.device)
        video = video_preprocess(frame_gen, height, width, self.dtype, self.device)
        if tiling_config is not None:
            return video_encoder.tiled_encode(video, tiling_config)
        return video_encoder(video)

    def _video_conditionings(
        self,
        images: list[ImageConditioningInput],
        driving_video_path: str,
        masks: torch.Tensor | None,
        height: int,
        width: int,
        num_frames: int,
        mode: DrivingMode,
        driving_strength: float,
        video_encoder: VideoEncoder,
        tiling_config: TilingConfig | None,
    ) -> list[ConditioningItem]:
        conditionings = combined_image_conditionings(
            images=images,
            height=height,
            width=width,
            video_encoder=video_encoder,
            dtype=self.dtype,
            device=self.device,
        )
        driving_latent = self._encode_driving_latent(
            driving_video_path, height, width, num_frames, video_encoder, tiling_config
        )
        conditionings.extend(
            build_scail_conditionings(driving_latent, masks, mode=mode, strength=driving_strength)
        )
        return conditionings

    @torch.inference_mode()
    def __call__(  # noqa: PLR0913
        self,
        prompt: str,
        seed: int,
        height: int,
        width: int,
        num_frames: int,
        frame_rate: float,
        driving_video_path: str,
        mask_path: str | None,
        images: list[ImageConditioningInput] | None = None,
        mode: DrivingMode = DrivingMode.ANIMATION,
        driving_strength: float = 1.0,
        enhance_prompt: bool = False,
        tiling_config: TilingConfig | None = None,
        stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
        stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
    ) -> tuple[Iterator[torch.Tensor], Audio | None]:
        assert_resolution(height=height, width=width, is_two_stage=True)
        images = images or []

        generator = torch.Generator(device=self.device).manual_seed(seed)
        noiser = GaussianNoiser(generator=generator)
        encode_tiling = TilingConfig.default()

        (ctx_p,) = self.prompt_encoder(
            [prompt],
            enhance_first_prompt=enhance_prompt,
            enhance_prompt_image=images[0][0] if len(images) > 0 else None,
            enhance_prompt_seed=seed,
        )
        video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding

        masks = load_masks(mask_path, self.device, self.dtype) if mask_path is not None else None

        def build_video_conditionings(out_height: int, out_width: int) -> list[ConditioningItem]:
            return self.image_conditioner(
                lambda enc: self._video_conditionings(
                    images=images,
                    driving_video_path=driving_video_path,
                    masks=masks,
                    height=out_height,
                    width=out_width,
                    num_frames=num_frames,
                    mode=mode,
                    driving_strength=driving_strength,
                    video_encoder=enc,
                    tiling_config=encode_tiling,
                )
            )

        # Stage 1: low resolution.
        stage_1_sigmas_t = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
        stage_1_conditionings = build_video_conditionings(height // 2, width // 2)
        video_state, _ = self.stage(
            denoiser=SimpleDenoiser(video_context, audio_context),
            sigmas=stage_1_sigmas_t,
            noiser=noiser,
            width=width // 2,
            height=height // 2,
            frames=num_frames,
            fps=frame_rate,
            video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
            audio=ModalitySpec(context=audio_context),
        )

        # Stage 2: upsample + refine.
        upscaled = self.upsampler(video_state.latent[:1])
        stage_2_sigmas_t = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
        stage_2_conditionings = build_video_conditionings(height, width)
        video_state, _ = self.stage(
            denoiser=SimpleDenoiser(video_context, audio_context),
            sigmas=stage_2_sigmas_t,
            noiser=noiser,
            width=width,
            height=height,
            frames=num_frames,
            fps=frame_rate,
            video=ModalitySpec(
                context=video_context,
                conditionings=stage_2_conditionings,
                noise_scale=stage_2_sigmas_t[0].item(),
                initial_latent=upscaled,
            ),
            audio=ModalitySpec(context=audio_context),
        )

        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
        return decoded_video, None


@torch.inference_mode()
def main() -> None:
    logging.basicConfig(level=logging.INFO)
    params = resolve_cli_params(distilled=True)
    parser = scail_animation_arg_parser(params=params)
    args = parser.parse_args()

    if not args.lora or len(args.lora) != 1:
        raise ValueError("SCAIL animation requires exactly one --lora (the SCAIL adapter).")

    pipeline = ScailAnimationPipeline(
        distilled_checkpoint_path=args.distilled_checkpoint_path,
        gemma_root=args.gemma_root,
        spatial_upsampler_path=args.spatial_upsampler_path,
        scail_lora=args.lora[0],
        quantization=args.quantization,
        compilation_config=args.compile,
        offload_mode=args.offload_mode,
    )
    tiling_config = TilingConfig.default()
    output_shape = VideoPixelShape(
        batch=1, frames=args.num_frames, width=args.width, height=args.height, fps=args.frame_rate
    )
    video_chunks_number = get_video_chunks_number(output_shape.frames, tiling_config)
    video, _audio = pipeline(
        prompt=args.prompt,
        seed=args.seed,
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        frame_rate=args.frame_rate,
        driving_video_path=args.driving_video,
        mask_path=args.mask_path,
        images=args.images if hasattr(args, "images") else [],
        mode=DrivingMode(args.mode),
        driving_strength=args.driving_strength,
        enhance_prompt=args.enhance_prompt,
        tiling_config=tiling_config,
    )
    encode_video(
        video=video,
        fps=int(args.frame_rate),
        audio=None,
        output_path=args.output_path,
        video_chunks_number=video_chunks_number,
    )


if __name__ == "__main__":
    main()