import PIL.Image
import torch

from diffusers import (
    AutoencoderKLLTX2Audio,
    AutoencoderKLLTX2Video,
    FlowMatchEulerDiscreteScheduler,
    LTX2ConditionPipeline,
    LTX2VideoTransformer3DModel,
)
from diffusers.modular_pipelines import ComponentSpec, InputParam, ModularPipelineBlocks, OutputParam, PipelineState
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
from diffusers.utils import load_image
from transformers import Gemma3ForConditionalGeneration, GemmaTokenizerFast


_LTX2_REPO = "diffusers/LTX-2.3-Diffusers"


class LTX2FirstLastFrameBlock(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Generates a video with LTX-2.3 conditioned on a first frame (`image`) and an optional "
            "last frame (`image_2`). Images can be PIL images, local paths, or URLs — path/URL strings "
            "are loaded automatically, so plain image inputs work without constructing "
            "`LTX2VideoCondition` objects."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler, pretrained_model_name_or_path=_LTX2_REPO, subfolder="scheduler"),
            ComponentSpec("vae", AutoencoderKLLTX2Video, pretrained_model_name_or_path=_LTX2_REPO, subfolder="vae"),
            ComponentSpec("audio_vae", AutoencoderKLLTX2Audio, pretrained_model_name_or_path=_LTX2_REPO, subfolder="audio_vae"),
            ComponentSpec("text_encoder", Gemma3ForConditionalGeneration, pretrained_model_name_or_path=_LTX2_REPO, subfolder="text_encoder"),
            ComponentSpec("tokenizer", GemmaTokenizerFast, pretrained_model_name_or_path=_LTX2_REPO, subfolder="tokenizer"),
            ComponentSpec("connectors", LTX2TextConnectors, pretrained_model_name_or_path=_LTX2_REPO, subfolder="connectors"),
            ComponentSpec("transformer", LTX2VideoTransformer3DModel, pretrained_model_name_or_path=_LTX2_REPO, subfolder="transformer"),
            ComponentSpec("vocoder", LTX2VocoderWithBWE, pretrained_model_name_or_path=_LTX2_REPO, subfolder="vocoder"),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam("prompt", type_hint=str, required=True, description="Text prompt describing the video."),
            InputParam(
                "image",
                type_hint="PIL.Image.Image | str",
                required=True,
                description="First frame to condition on. PIL image, local path, or URL.",
            ),
            InputParam(
                "image_2",
                type_hint="PIL.Image.Image | str | None",
                description="Optional last frame to condition on. PIL image, local path, or URL.",
            ),
            InputParam("negative_prompt", type_hint="str | None", description="Negative text prompt."),
            InputParam("height", type_hint=int, default=512, description="Output video height in pixels."),
            InputParam("width", type_hint=int, default=768, description="Output video width in pixels."),
            InputParam("num_frames", type_hint=int, default=121, description="Number of frames to generate."),
            InputParam("frame_rate", type_hint=float, default=24.0, description="Frames per second."),
            InputParam("num_inference_steps", type_hint=int, default=40, description="Number of denoising steps."),
            InputParam("guidance_scale", type_hint=float, default=4.0, description="Classifier-free guidance scale."),
            InputParam(
                "first_frame_strength", type_hint=float, default=1.0, description="Conditioning strength of `image`."
            ),
            InputParam(
                "last_frame_strength", type_hint=float, default=1.0, description="Conditioning strength of `image_2`."
            ),
            InputParam(
                "offload",
                type_hint=bool,
                default=True,
                description="Enable model CPU offload so the ~100GB of components fit a single 80GB GPU.",
            ),
            InputParam("generator", description="torch.Generator for deterministic sampling."),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("videos", type_hint="list[np.ndarray]", description="Generated video frames."),
            OutputParam("audio", type_hint=torch.Tensor, description="Generated audio waveform."),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        first = block_state.image
        if isinstance(first, str):
            first = load_image(first)
        conditions = [LTX2VideoCondition(frames=first, index=0, strength=block_state.first_frame_strength)]

        last = block_state.image_2
        if last is not None:
            if isinstance(last, str):
                last = load_image(last)
            conditions.append(LTX2VideoCondition(frames=last, index=-1, strength=block_state.last_frame_strength))

        pipeline = LTX2ConditionPipeline(
            scheduler=components.scheduler,
            vae=components.vae,
            audio_vae=components.audio_vae,
            text_encoder=components.text_encoder,
            tokenizer=components.tokenizer,
            connectors=components.connectors,
            transformer=components.transformer,
            vocoder=components.vocoder,
        )
        if block_state.offload:
            pipeline.enable_model_cpu_offload()

        result = pipeline(
            conditions=conditions,
            prompt=block_state.prompt,
            negative_prompt=block_state.negative_prompt,
            height=block_state.height,
            width=block_state.width,
            num_frames=block_state.num_frames,
            frame_rate=block_state.frame_rate,
            num_inference_steps=block_state.num_inference_steps,
            guidance_scale=block_state.guidance_scale,
            generator=block_state.generator,
            output_type="np",
        )
        # A list of per-frame ndarrays rather than PIL images: consumers (including
        # `diffusers-cli run`) can't distinguish a flat list of PIL frames from a batch
        # of independent images, but a list of HWC arrays is unambiguously a video.
        block_state.videos = list(result.frames[0])
        block_state.audio = result.audio

        self.set_block_state(state, block_state)
        return components, state