"""Video generation interface.

Supports:
- Google Gemini (Veo) via google-genai
- Volcengine Jimeng (即梦) via visual.volcengineapi.com (AK/SK signed)
"""

from __future__ import annotations

from io import BytesIO
import logging
from pathlib import Path
import time
from typing import Optional, Union, List, Tuple, Any

from google import genai
from google.genai import types
from PIL import Image

from ._rate_limit import rate_limited

logger = logging.getLogger(__name__)


DEFAULT_VIDEO_MODEL = "veo-3.1-generate-preview"
DEFAULT_JIMENG_VIDEO_MODEL = "jimeng_t2v_v30"


def _jimeng_frames_from_duration(duration_seconds: Optional[int], *, frames: Optional[int] = None) -> int:
    """Map generic duration to Jimeng frames (24*n+1)."""
    if frames is not None:
        if int(frames) not in (121, 241):
            raise ValueError("Jimeng frames must be one of: 121 (5s), 241 (10s)")
        return int(frames)
    if duration_seconds is None:
        return 121
    if int(duration_seconds) not in (5, 10):
        raise ValueError("Jimeng duration_seconds must be 5 or 10 (or pass frames=121/241 via extra_config)")
    return int(duration_seconds) * 24 + 1


def _guess_image_mime_type(image_bytes: bytes) -> str:
    try:
        img = Image.open(BytesIO(image_bytes))
        fmt = (img.format or "").upper()
    except Exception:
        return "image/png"

    if fmt == "JPEG":
        return "image/jpeg"
    if fmt == "WEBP":
        return "image/webp"
    if fmt == "GIF":
        return "image/gif"
    if fmt == "PNG":
        return "image/png"
    return "application/octet-stream"


def _build_image(input_image: Union[str, Path, bytes]) -> types.Image:
    if isinstance(input_image, (str, Path)):
        p = Path(input_image)
        if not p.exists():
            raise FileNotFoundError(f"Input image not found: {p}")
        return types.Image.from_file(location=str(p))

    if isinstance(input_image, (bytes, bytearray)):
        raw = bytes(input_image)
        return types.Image(image_bytes=raw, mime_type=_guess_image_mime_type(raw))

    raise TypeError(f"Unsupported input_image type: {type(input_image)!r}")


def _build_video(input_video: Union[str, Path, bytes, types.Video]) -> types.Video:
    if isinstance(input_video, types.Video):
        return input_video

    if isinstance(input_video, (str, Path)):
        s = str(input_video)
        p = Path(s)
        if p.exists():
            return types.Video.from_file(location=str(p))
        # Treat as URI if it isn't a readable local path.
        return types.Video(uri=s)

    if isinstance(input_video, (bytes, bytearray)):
        raw = bytes(input_video)
        # Default to MP4 if unknown.
        return types.Video(video_bytes=raw, mime_type="video/mp4")

    raise TypeError(f"Unsupported input_video type: {type(input_video)!r}")


def _parse_reference_type(
    reference_type: Optional[Union[str, types.VideoGenerationReferenceType]],
) -> Optional[types.VideoGenerationReferenceType]:
    if reference_type is None:
        return None
    if isinstance(reference_type, types.VideoGenerationReferenceType):
        return reference_type
    if not isinstance(reference_type, str):
        raise TypeError(f"Unsupported reference_type: {type(reference_type)!r}")

    key = reference_type.strip().upper()
    if key in ("ASSET",):
        return types.VideoGenerationReferenceType.ASSET
    if key in ("STYLE",):
        return types.VideoGenerationReferenceType.STYLE
    raise ValueError("reference_type must be one of: ASSET, STYLE")


ReferenceImageInput = Union[
    str,
    Path,
    bytes,
    types.VideoGenerationReferenceImage,
    Tuple[Union[str, Path, bytes], Union[str, types.VideoGenerationReferenceType]],
]


def _build_reference_images(
    reference_images: Optional[List[ReferenceImageInput]],
    *,
    reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
) -> Optional[List[types.VideoGenerationReferenceImage]]:
    if not reference_images:
        return None
    if len(reference_images) > 3:
        raise ValueError("reference_images supports at most 3 images")

    if reference_types is not None and len(reference_types) != len(reference_images):
        raise ValueError("reference_types must have the same length as reference_images")

    built: List[types.VideoGenerationReferenceImage] = []
    for idx, item in enumerate(reference_images):
        if isinstance(item, types.VideoGenerationReferenceImage):
            built.append(item)
            continue

        img: Union[str, Path, bytes]
        ref_type: Optional[Union[str, types.VideoGenerationReferenceType]] = None
        if isinstance(item, tuple) and len(item) == 2:
            img, ref_type = item
        else:
            img = item  # type: ignore[assignment]
            if reference_types is not None:
                ref_type = reference_types[idx]

        built.append(
            types.VideoGenerationReferenceImage(
                image=_build_image(img),
                reference_type=_parse_reference_type(ref_type),
            )
        )

    return built


def _merge_videos_config(
    base: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]],
    *,
    negative_prompt: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    resolution: Optional[str] = None,
    duration_seconds: Optional[int] = None,
    person_generation: Optional[str] = None,
    seed: Optional[int] = None,
    fps: Optional[int] = None,
    number_of_videos: Optional[int] = None,
    enhance_prompt: Optional[bool] = None,
    generate_audio: Optional[bool] = None,
    compression_quality: Optional[types.VideoCompressionQuality] = None,
    output_gcs_uri: Optional[str] = None,
    pubsub_topic: Optional[str] = None,
    last_frame: Optional[types.Image] = None,
    reference_images: Optional[List[types.VideoGenerationReferenceImage]] = None,
    extra: Optional[dict[str, Any]] = None,
) -> types.GenerateVideosConfig:
    data: dict[str, Any] = {}
    if base is not None:
        if isinstance(base, dict):
            data.update(base)
        else:
            data.update(base.model_dump(exclude_none=True))
    if extra:
        data.update(extra)

    # Explicit args override base.
    if negative_prompt is not None:
        data["negative_prompt"] = negative_prompt
    if aspect_ratio is not None:
        data["aspect_ratio"] = aspect_ratio
    if resolution is not None:
        data["resolution"] = resolution
    if duration_seconds is not None:
        data["duration_seconds"] = duration_seconds
    if person_generation is not None:
        data["person_generation"] = person_generation
    if seed is not None:
        data["seed"] = seed
    if fps is not None:
        data["fps"] = fps
    if number_of_videos is not None:
        data["number_of_videos"] = number_of_videos
    if enhance_prompt is not None:
        data["enhance_prompt"] = enhance_prompt
    if generate_audio is not None:
        data["generate_audio"] = generate_audio
    if compression_quality is not None:
        data["compression_quality"] = compression_quality
    if output_gcs_uri is not None:
        data["output_gcs_uri"] = output_gcs_uri
    if pubsub_topic is not None:
        data["pubsub_topic"] = pubsub_topic
    if last_frame is not None:
        data["last_frame"] = last_frame
    if reference_images is not None:
        data["reference_images"] = reference_images

    return types.GenerateVideosConfig.model_validate(data)


class VideoGenerator:
    """Wrapper for Veo video generation via google-genai."""

    def __init__(self, api_key: Optional[str] = None, *, provider: str = "gemini"):
        self.provider = (provider or "gemini").lower().strip()
        self.api_key = api_key
        if self.provider in ("gemini", "google", "veo"):
            self.client = genai.Client(api_key=api_key) if api_key else genai.Client()
            self.default_model = DEFAULT_VIDEO_MODEL
        elif self.provider in ("jimeng",):
            # Jimeng uses signed HTTP (AK/SK); no google-genai client needed.
            self.client = None  # type: ignore[assignment]
            self.default_model = DEFAULT_JIMENG_VIDEO_MODEL
        else:
            raise ValueError("provider must be one of: gemini, google, veo, jimeng")

    def generate_videos(
        self,
        prompt: str,
        output_path: Optional[Union[str, Path]] = None,
        *,
        model: Optional[str] = None,
        input_image: Optional[Union[str, Path, bytes]] = None,
        last_frame: Optional[Union[str, Path, bytes]] = None,
        reference_images: Optional[List[ReferenceImageInput]] = None,
        reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
        input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
        negative_prompt: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        duration_seconds: Optional[int] = None,
        person_generation: Optional[str] = None,
        seed: Optional[int] = None,
        fps: Optional[int] = None,
        number_of_videos: Optional[int] = None,
        enhance_prompt: Optional[bool] = None,
        generate_audio: Optional[bool] = None,
        compression_quality: Optional[types.VideoCompressionQuality] = None,
        output_gcs_uri: Optional[str] = None,
        pubsub_topic: Optional[str] = None,
        config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
        extra_config: Optional[dict[str, Any]] = None,
        poll_interval_s: float = 5.0,
        timeout_s: float = 600.0,
        save_description: bool = True,
    ) -> tuple[List[Path], Optional[Path]]:
        """Generate video(s) based on prompt (and optional inputs).

        Args:
            prompt: Text prompt for video generation.
            output_path: Output folder (or file path). If folder/no suffix, saves as
                generated_video.mp4 and generated_video_description.md by default.
            model: Optional model name. Default: veo-3.1-generate-preview
            input_image: Optional first-frame image (path or bytes).
            last_frame: Optional last frame image for interpolation (path or bytes).
            reference_images: Optional reference images (max 3). Each item can be:
                - image path/bytes
                - (image, reference_type) tuple, where type is ASSET or STYLE
                - types.VideoGenerationReferenceImage
            input_video: Optional input video for extension (path/uri/bytes/types.Video).
            negative_prompt: Optional negative prompt.
            aspect_ratio: e.g., '16:9', '9:16', '1:1'
            resolution: e.g., '720p', '1080p'
            duration_seconds: Desired duration in seconds.
            person_generation: Provider-specific policy string.
            seed: Deterministic seed.
            fps: Frames per second.
            number_of_videos: Number of videos to generate.
            enhance_prompt: Whether to enhance the prompt (if supported).
            generate_audio: Whether to generate audio (if supported).
            compression_quality: Optional compression quality enum.
            output_gcs_uri: Optional GCS output URI.
            pubsub_topic: Optional Pub/Sub topic for async notifications.
            config: Optional base GenerateVideosConfig (merged with explicit args).
            extra_config: Optional dict merged into config (advanced / forward-compatible).
            poll_interval_s: Poll interval for long-running operation.
            timeout_s: Overall timeout for the generation request.
            save_description: Whether to save a markdown description file.

        Returns:
            (video_paths, description_path)
        """
        if self.provider == "jimeng":
            return self._generate_videos_jimeng(
                prompt,
                output_path,
                model=model,
                negative_prompt=negative_prompt,
                aspect_ratio=aspect_ratio,
                duration_seconds=duration_seconds,
                seed=seed,
                number_of_videos=number_of_videos,
                config=config,
                extra_config=extra_config,
                poll_interval_s=poll_interval_s,
                timeout_s=timeout_s,
                save_description=save_description,
                # unsupported (keep explicit for better error messages)
                input_image=input_image,
                last_frame=last_frame,
                reference_images=reference_images,
                reference_types=reference_types,
                input_video=input_video,
                resolution=resolution,
                person_generation=person_generation,
                fps=fps,
                enhance_prompt=enhance_prompt,
                generate_audio=generate_audio,
                compression_quality=compression_quality,
                output_gcs_uri=output_gcs_uri,
                pubsub_topic=pubsub_topic,
            )

        if not prompt:
            raise ValueError("prompt cannot be empty")
        if poll_interval_s <= 0:
            raise ValueError("poll_interval_s must be > 0")
        if timeout_s <= 0:
            raise ValueError("timeout_s must be > 0")

        use_model = model or self.default_model

        # Resolve output paths.
        if output_path is None:
            output_dir = Path(".")
            base_name = "generated_video"
            suffix = ".mp4"
        else:
            output_path = Path(output_path)
            if output_path.exists() and output_path.is_dir():
                output_dir = output_path
                base_name = "generated_video"
                suffix = ".mp4"
            elif not output_path.suffix:
                # Treat as folder/name combination.
                if output_path.parent != Path("."):
                    output_dir = output_path.parent
                    base_name = output_path.name
                else:
                    output_dir = Path(".")
                    base_name = output_path.name
                suffix = ".mp4"
            else:
                output_dir = output_path.parent if output_path.parent != Path(".") else Path(".")
                base_name = output_path.stem
                suffix = output_path.suffix

        output_dir.mkdir(parents=True, exist_ok=True)
        description_path = output_dir / f"{base_name}_description.md" if save_description else None

        image_obj = _build_image(input_image) if input_image is not None else None
        last_frame_obj = _build_image(last_frame) if last_frame is not None else None
        reference_images_obj = _build_reference_images(reference_images, reference_types=reference_types)
        video_obj = _build_video(input_video) if input_video is not None else None

        cfg = _merge_videos_config(
            config,
            negative_prompt=negative_prompt,
            aspect_ratio=aspect_ratio,
            resolution=resolution,
            duration_seconds=duration_seconds,
            person_generation=person_generation,
            seed=seed,
            fps=fps,
            number_of_videos=number_of_videos,
            enhance_prompt=enhance_prompt,
            generate_audio=generate_audio,
            compression_quality=compression_quality,
            output_gcs_uri=output_gcs_uri,
            pubsub_topic=pubsub_topic,
            last_frame=last_frame_obj,
            reference_images=reference_images_obj,
            extra=extra_config,
        )

        logger.info("Generating video with model=%s", use_model)
        operation = self.client.models.generate_videos(
            model=use_model,
            prompt=prompt,
            image=image_obj,
            video=video_obj,
            config=cfg,
        )

        start = time.time()
        while not getattr(operation, "done", False):
            if (time.time() - start) >= timeout_s:
                raise TimeoutError(f"Video generation timed out after {timeout_s:.0f}s")
            time.sleep(poll_interval_s)
            operation = self.client.operations.get(operation)

        response = getattr(operation, "result", None) or getattr(operation, "response", None)
        generated = getattr(response, "generated_videos", None) if response is not None else None
        if not generated:
            filtered = getattr(response, "rai_media_filtered_count", None) if response is not None else None
            if filtered:
                raise RuntimeError(f"Video generation returned no videos (filtered_count={filtered})")
            raise RuntimeError("Video generation returned no videos")

        # Download and persist all returned videos.
        video_paths: List[Path] = []
        for idx, item in enumerate(generated):
            v = getattr(item, "video", None)
            if v is None:
                continue
            data = self.client.files.download(file=v)
            path = output_dir / f"{base_name}{'' if len(generated) == 1 else f'_{idx}'}{suffix}"
            path.write_bytes(data)
            video_paths.append(path)

        if not video_paths:
            raise RuntimeError("Video generation returned no downloadable video payloads")

        if save_description and description_path:
            content = "# Video Generation Description\n\n"
            content += "**Generated Video(s)**:\n\n"
            for p in video_paths:
                content += f"- {p.name}\n"
            content += "\n"
            content += "## Prompt\n\n"
            content += f"{prompt}\n\n"
            if input_image is not None and isinstance(input_image, (str, Path)):
                content += "## Input Image\n\n"
                content += f"- {Path(input_image).name}\n\n"
            if last_frame is not None and isinstance(last_frame, (str, Path)):
                content += "## Last Frame\n\n"
                content += f"- {Path(last_frame).name}\n\n"
            if reference_images:
                content += "## Reference Images\n\n"
                for r in reference_images:
                    if isinstance(r, tuple) and len(r) == 2:
                        r = r[0]
                    if isinstance(r, (str, Path)):
                        content += f"- {Path(r).name}\n"
                    else:
                        content += "- (bytes)\n"
                content += "\n"
            if input_video is not None:
                content += "## Input Video (Extend)\n\n"
                if isinstance(input_video, (str, Path)):
                    content += f"- {str(input_video)}\n\n"
                else:
                    content += "- (bytes)\n\n"
            content += "## Generation Details\n\n"
            content += f"- **Model**: {use_model}\n"
            content += f"- **Config**: {cfg.model_dump(exclude_none=True)}\n"
            content += f"- **Timestamp**: {time.strftime('%Y-%m-%d %H:%M:%S')}\n"
            description_path.write_text(content, encoding="utf-8")

        return video_paths, description_path

    def _generate_videos_jimeng(
        self,
        prompt: str,
        output_path: Optional[Union[str, Path]] = None,
        *,
        model: Optional[str] = None,
        negative_prompt: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        duration_seconds: Optional[int] = None,
        seed: Optional[int] = None,
        number_of_videos: Optional[int] = None,
        config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
        extra_config: Optional[dict[str, Any]] = None,
        poll_interval_s: float = 5.0,
        timeout_s: float = 600.0,
        save_description: bool = True,
        # unsupported args (for clearer validation)
        input_image: Optional[Union[str, Path, bytes]] = None,
        last_frame: Optional[Union[str, Path, bytes]] = None,
        reference_images: Optional[List[ReferenceImageInput]] = None,
        reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
        input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
        resolution: Optional[str] = None,
        person_generation: Optional[str] = None,
        fps: Optional[int] = None,
        enhance_prompt: Optional[bool] = None,
        generate_audio: Optional[bool] = None,
        compression_quality: Optional[types.VideoCompressionQuality] = None,
        output_gcs_uri: Optional[str] = None,
        pubsub_topic: Optional[str] = None,
    ) -> tuple[List[Path], Optional[Path]]:
        from .jimeng_video import JimengVideoClient, DEFAULT_JIMENG_VIDEO_MODEL

        if not prompt:
            raise ValueError("prompt cannot be empty")
        if poll_interval_s <= 0:
            raise ValueError("poll_interval_s must be > 0")
        if timeout_s <= 0:
            raise ValueError("timeout_s must be > 0")

        unsupported: dict[str, Any] = {
            "input_image": input_image,
            "last_frame": last_frame,
            "reference_images": reference_images,
            "reference_types": reference_types,
            "input_video": input_video,
            "negative_prompt": negative_prompt,
            "resolution": resolution,
            "person_generation": person_generation,
            "fps": fps,
            "enhance_prompt": enhance_prompt,
            "generate_audio": generate_audio,
            "compression_quality": compression_quality,
            "output_gcs_uri": output_gcs_uri,
            "pubsub_topic": pubsub_topic,
            "config": config,
        }
        bad = [k for k, v in unsupported.items() if v is not None]
        if bad:
            raise ValueError(f"Jimeng provider does not support: {', '.join(bad)}")

        if number_of_videos is not None and int(number_of_videos) != 1:
            raise ValueError("Jimeng provider does not support number_of_videos != 1")

        use_model = model or DEFAULT_JIMENG_VIDEO_MODEL
        frames = _jimeng_frames_from_duration(duration_seconds, frames=(extra_config or {}).get("frames") if extra_config else None)
        use_aspect_ratio = aspect_ratio or "16:9"
        use_seed = int(seed) if seed is not None else -1

        # Resolve output paths (same behavior as Veo path handling).
        if output_path is None:
            output_dir = Path(".")
            base_name = "generated_video"
            suffix = ".mp4"
        else:
            output_path = Path(output_path)
            if output_path.exists() and output_path.is_dir():
                output_dir = output_path
                base_name = "generated_video"
                suffix = ".mp4"
            elif not output_path.suffix:
                if output_path.parent != Path("."):
                    output_dir = output_path.parent
                    base_name = output_path.name
                else:
                    output_dir = Path(".")
                    base_name = output_path.name
                suffix = ".mp4"
            else:
                output_dir = output_path.parent if output_path.parent != Path(".") else Path(".")
                base_name = output_path.stem
                suffix = output_path.suffix

        output_dir.mkdir(parents=True, exist_ok=True)
        description_path = output_dir / f"{base_name}_description.md" if save_description else None

        client = JimengVideoClient()
        logger.info("Submitting Jimeng video task model=%s frames=%s aspect_ratio=%s", use_model, frames, use_aspect_ratio)
        task_id = client.submit_t2v(
            prompt=prompt,
            model=use_model,
            seed=use_seed,
            frames=frames,
            aspect_ratio=use_aspect_ratio,
        )
        video_url = client.wait_for_video_url(
            model=use_model,
            task_id=task_id,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
        )
        video_bytes = client.download_video_bytes(video_url=video_url)

        out_path = output_dir / f"{base_name}{suffix}"
        out_path.write_bytes(video_bytes)

        if description_path is not None:
            content = "# Video Generation Description\n\n"
            content += f"**Provider**: jimeng\n\n"
            content += f"**Model (req_key)**: {use_model}\n\n"
            content += f"**Task ID**: {task_id}\n\n"
            content += f"**Aspect Ratio**: {use_aspect_ratio}\n\n"
            content += f"**Frames**: {frames}\n\n"
            content += f"## Prompt\n\n{prompt}\n"
            description_path.write_text(content, encoding="utf-8")

        return [out_path], description_path

    def generate_video(
        self,
        prompt: str,
        output_path: Optional[Union[str, Path]] = None,
        *,
        model: Optional[str] = None,
        input_image: Optional[Union[str, Path, bytes]] = None,
        last_frame: Optional[Union[str, Path, bytes]] = None,
        reference_images: Optional[List[ReferenceImageInput]] = None,
        reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
        input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
        negative_prompt: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        resolution: Optional[str] = None,
        duration_seconds: Optional[int] = None,
        person_generation: Optional[str] = None,
        seed: Optional[int] = None,
        fps: Optional[int] = None,
        enhance_prompt: Optional[bool] = None,
        generate_audio: Optional[bool] = None,
        compression_quality: Optional[types.VideoCompressionQuality] = None,
        output_gcs_uri: Optional[str] = None,
        pubsub_topic: Optional[str] = None,
        config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
        extra_config: Optional[dict[str, Any]] = None,
        poll_interval_s: float = 5.0,
        timeout_s: float = 600.0,
        save_description: bool = True,
    ) -> tuple[Path, Optional[Path]]:
        """Generate a single video (returns the first one if multiple are produced)."""
        paths, desc = self.generate_videos(
            prompt,
            output_path,
            model=model,
            input_image=input_image,
            last_frame=last_frame,
            reference_images=reference_images,
            reference_types=reference_types,
            input_video=input_video,
            negative_prompt=negative_prompt,
            aspect_ratio=aspect_ratio,
            resolution=resolution,
            duration_seconds=duration_seconds,
            person_generation=person_generation,
            seed=seed,
            fps=fps,
            number_of_videos=1,
            enhance_prompt=enhance_prompt,
            generate_audio=generate_audio,
            compression_quality=compression_quality,
            output_gcs_uri=output_gcs_uri,
            pubsub_topic=pubsub_topic,
            config=config,
            extra_config=extra_config,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
            save_description=save_description,
        )
        return paths[0], desc


@rate_limited(max_calls=4, period=60.0, key="api_sdk.generate_video")
def generate_video(
    prompt: str,
    output_path: Optional[Union[str, Path]] = None,
    *,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_image: Optional[Union[str, Path, bytes]] = None,
    last_frame: Optional[Union[str, Path, bytes]] = None,
    reference_images: Optional[List[ReferenceImageInput]] = None,
    reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
    input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
    negative_prompt: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    resolution: Optional[str] = None,
    duration_seconds: Optional[int] = None,
    person_generation: Optional[str] = None,
    seed: Optional[int] = None,
    fps: Optional[int] = None,
    enhance_prompt: Optional[bool] = None,
    generate_audio: Optional[bool] = None,
    compression_quality: Optional[types.VideoCompressionQuality] = None,
    output_gcs_uri: Optional[str] = None,
    pubsub_topic: Optional[str] = None,
    config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    poll_interval_s: float = 5.0,
    timeout_s: float = 600.0,
    save_description: bool = True,
) -> tuple[Path, Optional[Path]]:
    """Convenience function for video generation (file output)."""
    return VideoGenerator(api_key=api_key, provider=provider).generate_video(
        prompt,
        output_path,
        model=model,
        input_image=input_image,
        last_frame=last_frame,
        reference_images=reference_images,
        reference_types=reference_types,
        input_video=input_video,
        negative_prompt=negative_prompt,
        aspect_ratio=aspect_ratio,
        resolution=resolution,
        duration_seconds=duration_seconds,
        person_generation=person_generation,
        seed=seed,
        fps=fps,
        enhance_prompt=enhance_prompt,
        generate_audio=generate_audio,
        compression_quality=compression_quality,
        output_gcs_uri=output_gcs_uri,
        pubsub_topic=pubsub_topic,
        config=config,
        extra_config=extra_config,
        poll_interval_s=poll_interval_s,
        timeout_s=timeout_s,
        save_description=save_description,
    )


@rate_limited(max_calls=4, period=60.0, key="api_sdk.generate_videos")
def generate_videos(
    prompt: str,
    output_path: Optional[Union[str, Path]] = None,
    *,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_image: Optional[Union[str, Path, bytes]] = None,
    last_frame: Optional[Union[str, Path, bytes]] = None,
    reference_images: Optional[List[ReferenceImageInput]] = None,
    reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
    input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
    negative_prompt: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    resolution: Optional[str] = None,
    duration_seconds: Optional[int] = None,
    person_generation: Optional[str] = None,
    seed: Optional[int] = None,
    fps: Optional[int] = None,
    number_of_videos: Optional[int] = None,
    enhance_prompt: Optional[bool] = None,
    generate_audio: Optional[bool] = None,
    compression_quality: Optional[types.VideoCompressionQuality] = None,
    output_gcs_uri: Optional[str] = None,
    pubsub_topic: Optional[str] = None,
    config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    poll_interval_s: float = 5.0,
    timeout_s: float = 600.0,
    save_description: bool = True,
) -> tuple[List[Path], Optional[Path]]:
    """Generate multiple videos (file output)."""
    return VideoGenerator(api_key=api_key, provider=provider).generate_videos(
        prompt,
        output_path,
        model=model,
        input_image=input_image,
        last_frame=last_frame,
        reference_images=reference_images,
        reference_types=reference_types,
        input_video=input_video,
        negative_prompt=negative_prompt,
        aspect_ratio=aspect_ratio,
        resolution=resolution,
        duration_seconds=duration_seconds,
        person_generation=person_generation,
        seed=seed,
        fps=fps,
        number_of_videos=number_of_videos,
        enhance_prompt=enhance_prompt,
        generate_audio=generate_audio,
        compression_quality=compression_quality,
        output_gcs_uri=output_gcs_uri,
        pubsub_topic=pubsub_topic,
        config=config,
        extra_config=extra_config,
        poll_interval_s=poll_interval_s,
        timeout_s=timeout_s,
        save_description=save_description,
    )


@rate_limited(max_calls=4, period=60.0, key="api_sdk.generate_video_bytes")
def generate_video_bytes(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_image: Optional[Union[str, Path, bytes]] = None,
    last_frame: Optional[Union[str, Path, bytes]] = None,
    reference_images: Optional[List[ReferenceImageInput]] = None,
    reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
    input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
    negative_prompt: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    resolution: Optional[str] = None,
    duration_seconds: Optional[int] = None,
    person_generation: Optional[str] = None,
    seed: Optional[int] = None,
    fps: Optional[int] = None,
    enhance_prompt: Optional[bool] = None,
    generate_audio: Optional[bool] = None,
    compression_quality: Optional[types.VideoCompressionQuality] = None,
    output_gcs_uri: Optional[str] = None,
    pubsub_topic: Optional[str] = None,
    config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    poll_interval_s: float = 5.0,
    timeout_s: float = 600.0,
) -> bytes:
    """Generate a video and return raw bytes (no file IO)."""
    if (provider or "gemini").lower().strip() == "jimeng":
        from .jimeng_video import JimengVideoClient, DEFAULT_JIMENG_VIDEO_MODEL

        use_model = model or DEFAULT_JIMENG_VIDEO_MODEL
        if any(x is not None for x in (input_image, last_frame, reference_images, reference_types, input_video, negative_prompt, resolution, person_generation, fps, enhance_prompt, generate_audio, compression_quality, output_gcs_uri, pubsub_topic, config)):
            raise ValueError("Jimeng provider in generate_video_bytes only supports: prompt, model(req_key), aspect_ratio, duration_seconds, seed, extra_config(frames), poll_interval_s, timeout_s")
        frames = _jimeng_frames_from_duration(duration_seconds, frames=(extra_config or {}).get("frames") if extra_config else None)
        use_aspect_ratio = aspect_ratio or "16:9"
        use_seed = int(seed) if seed is not None else -1

        client = JimengVideoClient()
        task_id = client.submit_t2v(prompt=prompt, model=use_model, seed=use_seed, frames=frames, aspect_ratio=use_aspect_ratio)
        video_url = client.wait_for_video_url(model=use_model, task_id=task_id, poll_interval_s=poll_interval_s, timeout_s=timeout_s)
        return client.download_video_bytes(video_url=video_url)

    gen = VideoGenerator(api_key=api_key, provider=provider)
    use_model = model or gen.default_model
    image_obj = _build_image(input_image) if input_image is not None else None
    last_frame_obj = _build_image(last_frame) if last_frame is not None else None
    reference_images_obj = _build_reference_images(reference_images, reference_types=reference_types)
    video_obj = _build_video(input_video) if input_video is not None else None

    cfg = _merge_videos_config(
        config,
        negative_prompt=negative_prompt,
        aspect_ratio=aspect_ratio,
        resolution=resolution,
        duration_seconds=duration_seconds,
        person_generation=person_generation,
        seed=seed,
        fps=fps,
        number_of_videos=1,
        enhance_prompt=enhance_prompt,
        generate_audio=generate_audio,
        compression_quality=compression_quality,
        output_gcs_uri=output_gcs_uri,
        pubsub_topic=pubsub_topic,
        last_frame=last_frame_obj,
        reference_images=reference_images_obj,
        extra=extra_config,
    )

    operation = gen.client.models.generate_videos(model=use_model, prompt=prompt, image=image_obj, video=video_obj, config=cfg)
    start = time.time()
    while not getattr(operation, "done", False):
        if (time.time() - start) >= timeout_s:
            raise TimeoutError(f"Video generation timed out after {timeout_s:.0f}s")
        time.sleep(poll_interval_s)
        operation = gen.client.operations.get(operation)

    response = getattr(operation, "result", None) or getattr(operation, "response", None)
    generated = getattr(response, "generated_videos", None) if response is not None else None
    if not generated:
        filtered = getattr(response, "rai_media_filtered_count", None) if response is not None else None
        if filtered:
            raise RuntimeError(f"Video generation returned no videos (filtered_count={filtered})")
        raise RuntimeError("Video generation returned no videos")
    video_obj = generated[0].video
    if video_obj is None:
        raise RuntimeError("Video generation returned an empty video object")
    return gen.client.files.download(file=video_obj)


@rate_limited(max_calls=4, period=60.0, key="api_sdk.generate_videos_bytes")
def generate_videos_bytes(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_image: Optional[Union[str, Path, bytes]] = None,
    last_frame: Optional[Union[str, Path, bytes]] = None,
    reference_images: Optional[List[ReferenceImageInput]] = None,
    reference_types: Optional[List[Union[str, types.VideoGenerationReferenceType]]] = None,
    input_video: Optional[Union[str, Path, bytes, types.Video]] = None,
    negative_prompt: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    resolution: Optional[str] = None,
    duration_seconds: Optional[int] = None,
    person_generation: Optional[str] = None,
    seed: Optional[int] = None,
    fps: Optional[int] = None,
    number_of_videos: Optional[int] = None,
    enhance_prompt: Optional[bool] = None,
    generate_audio: Optional[bool] = None,
    compression_quality: Optional[types.VideoCompressionQuality] = None,
    output_gcs_uri: Optional[str] = None,
    pubsub_topic: Optional[str] = None,
    config: Optional[Union[types.GenerateVideosConfig, types.GenerateVideosConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    poll_interval_s: float = 5.0,
    timeout_s: float = 600.0,
) -> List[bytes]:
    """Generate multiple videos and return their raw bytes."""
    if (provider or "gemini").lower().strip() == "jimeng":
        if number_of_videos is not None and int(number_of_videos) != 1:
            raise ValueError("Jimeng provider does not support number_of_videos != 1")
        return [generate_video_bytes(
            prompt,
            api_key=api_key,
            provider=provider,
            model=model,
            input_image=input_image,
            last_frame=last_frame,
            reference_images=reference_images,
            reference_types=reference_types,
            input_video=input_video,
            negative_prompt=negative_prompt,
            aspect_ratio=aspect_ratio,
            resolution=resolution,
            duration_seconds=duration_seconds,
            person_generation=person_generation,
            seed=seed,
            fps=fps,
            enhance_prompt=enhance_prompt,
            generate_audio=generate_audio,
            compression_quality=compression_quality,
            output_gcs_uri=output_gcs_uri,
            pubsub_topic=pubsub_topic,
            config=config,
            extra_config=extra_config,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
        )]

    gen = VideoGenerator(api_key=api_key, provider=provider)
    use_model = model or gen.default_model
    image_obj = _build_image(input_image) if input_image is not None else None
    last_frame_obj = _build_image(last_frame) if last_frame is not None else None
    reference_images_obj = _build_reference_images(reference_images, reference_types=reference_types)
    video_obj = _build_video(input_video) if input_video is not None else None

    cfg = _merge_videos_config(
        config,
        negative_prompt=negative_prompt,
        aspect_ratio=aspect_ratio,
        resolution=resolution,
        duration_seconds=duration_seconds,
        person_generation=person_generation,
        seed=seed,
        fps=fps,
        number_of_videos=number_of_videos,
        enhance_prompt=enhance_prompt,
        generate_audio=generate_audio,
        compression_quality=compression_quality,
        output_gcs_uri=output_gcs_uri,
        pubsub_topic=pubsub_topic,
        last_frame=last_frame_obj,
        reference_images=reference_images_obj,
        extra=extra_config,
    )

    operation = gen.client.models.generate_videos(model=use_model, prompt=prompt, image=image_obj, video=video_obj, config=cfg)
    start = time.time()
    while not getattr(operation, "done", False):
        if (time.time() - start) >= timeout_s:
            raise TimeoutError(f"Video generation timed out after {timeout_s:.0f}s")
        time.sleep(poll_interval_s)
        operation = gen.client.operations.get(operation)

    response = getattr(operation, "result", None) or getattr(operation, "response", None)
    generated = getattr(response, "generated_videos", None) if response is not None else None
    if not generated:
        filtered = getattr(response, "rai_media_filtered_count", None) if response is not None else None
        if filtered:
            raise RuntimeError(f"Video generation returned no videos (filtered_count={filtered})")
        raise RuntimeError("Video generation returned no videos")

    out: List[bytes] = []
    for item in generated:
        v = getattr(item, "video", None)
        if v is None:
            continue
        out.append(gen.client.files.download(file=v))
    if not out:
        raise RuntimeError("Video generation returned no downloadable video payloads")
    return out
