from __future__ import annotations

from typing import Any

import torch
from diffusers import DiffusionPipeline

from .base import ImageToVideoGenerator


class DiffusersImageToVideoAdapter(ImageToVideoGenerator):
    def __init__(self, model_id: str, device: str = "cuda", torch_dtype: Any | None = None, pipeline: Any | None = None, **kwargs) -> None:
        if pipeline is None:
            dtype = torch_dtype
            if dtype is None and str(device).startswith("cuda"):
                dtype = torch.bfloat16
            pipeline_cls = kwargs.pop("pipeline_cls", None) or _default_i2v_pipeline_cls(model_id)
            pipeline = pipeline_cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
        self.pipe = pipeline
        self.device = device
        if hasattr(self.pipe, "to"):
            self.pipe.to(device)
        if hasattr(self.pipe, "set_progress_bar_config"):
            self.pipe.set_progress_bar_config(disable=True)

    def generate(self, image: Any, prompt: str, *, generator: Any | None = None, **kwargs) -> Any:
        kwargs.setdefault("output_type", "pt")
        try:
            out = self.pipe(image=image, prompt=prompt, generator=generator, **kwargs)
        except TypeError as exc:
            raise TypeError(
                "The configured image-to-video model does not support the default "
                "`image=..., prompt=..., generator=..., **kwargs` signature. "
                "Pass a custom ImageToVideoGenerator adapter."
            ) from exc
        return self._normalize_output(out)

    def generate_batch(self, images, prompts, *, generators=None, **kwargs):
        images = list(images)
        prompts = list(prompts)
        if not prompts:
            return []
        kwargs.setdefault("output_type", "pt")
        try:
            out = self.pipe(image=images, prompt=prompts, generator=generators, **kwargs)
            frames = getattr(out, "frames", None)
            if frames is None:
                frames = getattr(out, "videos", None)
            if torch.is_tensor(frames) and frames.ndim == 5 and frames.shape[0] == len(prompts):
                return [frames[i] for i in range(len(prompts))]
            if isinstance(frames, (list, tuple)) and len(frames) == len(prompts):
                return list(frames)
        except (RuntimeError, TypeError, ValueError):
            pass
        # Fall back to per-item generation (e.g. model can't batch, or OOM).
        gens = generators if isinstance(generators, (list, tuple)) else [generators] * len(prompts)
        return [self.generate(img, p, generator=g, **kwargs) for img, p, g in zip(images, prompts, gens)]

    @staticmethod
    def _normalize_output(out: Any) -> Any:
        frames = getattr(out, "frames", None)
        if frames is None:
            frames = getattr(out, "videos", None)
        if frames is None:
            return out
        if torch.is_tensor(frames):
            return frames[0] if frames.ndim == 5 else frames
        if isinstance(frames, (list, tuple)) and frames:
            return frames[0]
        return frames


def _default_i2v_pipeline_cls(model_id: str):
    if model_id == "Lightricks/LTX-Video":
        from diffusers import LTXImageToVideoPipeline

        return LTXImageToVideoPipeline
    return DiffusionPipeline