# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Shared helper utilities for OpenAI-compatible video generation API.
"""

from __future__ import annotations

import base64
import binascii
from collections import deque
from io import BytesIO
from typing import Any, Literal

import httpx
import numpy as np
import torch
from PIL import Image, UnidentifiedImageError

from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
from vllm_omni.entrypoints.openai.protocol.videos import (
    FileImageReference,
    FileVideoReference,
    ImageReference,
    UrlImageReference,
    UrlVideoReference,
    VideoReference,
)


class VideoFrames(list[Image.Image]):
    """Decoded video frames plus source metadata."""

    def __init__(self, frames: list[Image.Image] | None = None, *, fps: float | None = None) -> None:
        super().__init__(frames or [])
        self.fps = fps
        self.frame_rate = fps


def positive_float(value: Any) -> float | None:
    if value is None:
        return None
    if hasattr(value, "item") and not isinstance(value, (bytes, str)):
        value = value.item()
    try:
        result = float(value)
    except (TypeError, ValueError):
        return None
    if result <= 0:
        return None
    return result


def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
    try:
        return Image.open(BytesIO(image_bytes)).convert("RGB")
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc


def _decode_video_bytes(
    video_bytes: bytes,
    *,
    source: str,
    max_frames: int | None = None,
    keep: Literal["first", "last"] = "first",
) -> VideoFrames:
    try:
        import av
    except ImportError as exc:  # pragma: no cover - av is a serving dependency via media_utils
        raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc

    if keep not in {"first", "last"}:
        raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.")
    if max_frames is not None and max_frames <= 0:
        raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.")

    frames: list[Image.Image] = []
    tail_frames: deque[Image.Image] | None = (
        deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None
    )
    fps: float | None = None
    try:
        with av.open(BytesIO(video_bytes)) as container:
            video_stream = container.streams.video[0] if container.streams.video else None
            if video_stream is not None:
                fps = (
                    positive_float(getattr(video_stream, "average_rate", None))
                    or positive_float(getattr(video_stream, "base_rate", None))
                    or positive_float(getattr(video_stream, "guessed_rate", None))
                )
            for frame in container.decode(video=0):
                image = frame.to_image().convert("RGB")
                if tail_frames is not None:
                    tail_frames.append(image)
                else:
                    frames.append(image)
                if keep == "first" and max_frames is not None and len(frames) >= max_frames:
                    break
    except Exception as exc:
        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc

    if tail_frames is not None:
        frames = list(tail_frames)
    if not frames:
        raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.")
    return VideoFrames(frames, fps=fps)


def _decode_media_bytes(
    media_bytes: bytes,
    *,
    source: str,
    max_video_frames: int | None = None,
    video_keep: Literal["first", "last"] = "first",
) -> Image.Image | VideoFrames:
    try:
        return _decode_image_bytes(media_bytes, source=source)
    except InvalidInputReferenceError:
        try:
            return _decode_video_bytes(
                media_bytes,
                source=source,
                max_frames=max_video_frames,
                keep=video_keep,
            )
        except InvalidInputReferenceError as video_exc:
            raise InvalidInputReferenceError(
                f"Invalid {source}: provided content is not a valid image or video."
            ) from video_exc


def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
    if input_reference:
        if input_reference.startswith("data:image"):
            _, b64_data = input_reference.split(",", 1)
        else:
            b64_data = input_reference

        try:
            image_bytes = base64.b64decode(b64_data)
        except (binascii.Error, ValueError) as exc:  # pragma: no cover - malformed base64
            raise InvalidInputReferenceError(f"Invalid {source}: image data is not valid base64.") from exc
        return _decode_image_bytes(image_bytes, source=source)
    raise InvalidInputReferenceError(f"Invalid {source}: image data is empty.")


async def decode_image_url(image_url: str) -> Image.Image:
    if image_url.startswith("data:image"):
        return _decode_base64_image(image_url, source="image_reference.image_url")

    if image_url.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60) as client:
            try:
                response = await client.get(image_url)
                response.raise_for_status()
            except httpx.HTTPError as exc:
                raise InvalidInputReferenceError(
                    "Invalid image_reference.image_url: failed to download image."
                ) from exc
        return _decode_image_bytes(response.content, source="image_reference.image_url")

    raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")


def _decode_base64_video(
    video_reference: str,
    *,
    source: str,
    max_frames: int | None = None,
    keep: Literal["first", "last"] = "first",
) -> VideoFrames:
    if video_reference:
        if video_reference.startswith("data:video"):
            _, b64_data = video_reference.split(",", 1)
        else:
            b64_data = video_reference

        try:
            video_bytes = base64.b64decode(b64_data)
        except (binascii.Error, ValueError) as exc:  # pragma: no cover - malformed base64
            raise InvalidInputReferenceError(f"Invalid {source}: video data is not valid base64.") from exc
        return _decode_video_bytes(video_bytes, source=source, max_frames=max_frames, keep=keep)
    raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")


async def decode_video_url(
    video_url: str,
    *,
    max_frames: int | None = None,
    keep: Literal["first", "last"] = "first",
) -> VideoFrames:
    if video_url.startswith("data:video"):
        return _decode_base64_video(
            video_url,
            source="video_reference.video_url",
            max_frames=max_frames,
            keep=keep,
        )

    if video_url.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60) as client:
            try:
                response = await client.get(video_url)
                response.raise_for_status()
            except httpx.HTTPError as exc:
                raise InvalidInputReferenceError(
                    "Invalid video_reference.video_url: failed to download video."
                ) from exc
        return _decode_video_bytes(
            response.content,
            source="video_reference.video_url",
            max_frames=max_frames,
            keep=keep,
        )

    raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")


async def decode_audio_url(audio_url: str) -> str:
    """Decode an audio URL or data-URL to a temporary file path."""
    import tempfile

    audio_bytes: bytes | None = None

    if audio_url.startswith("data:audio"):
        _, b64_data = audio_url.split(",", 1)
        try:
            audio_bytes = base64.b64decode(b64_data)
        except (binascii.Error, ValueError) as exc:
            raise InvalidInputReferenceError(
                "Invalid audio_reference.audio_url: audio data is not valid base64."
            ) from exc
    elif audio_url.startswith(("http://", "https://")):
        async with httpx.AsyncClient(timeout=60) as client:
            try:
                response = await client.get(audio_url)
                response.raise_for_status()
            except httpx.HTTPError as exc:
                raise InvalidInputReferenceError(
                    "Invalid audio_reference.audio_url: failed to download audio."
                ) from exc
        audio_bytes = response.content
    else:
        raise InvalidInputReferenceError("Invalid audio_reference.audio_url: must be an http(s) URL or data URL.")

    if not audio_bytes:
        raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.")

    suffix = ".wav"
    if audio_url.startswith("data:audio/"):
        mime = audio_url.split(";")[0].removeprefix("data:")
        ext = mime.split("/")[-1]
        if ext in ("mpeg", "mp3"):
            suffix = ".mp3"
        elif ext == "wav":
            suffix = ".wav"
        elif ext.isalnum() and len(ext) <= 8:
            suffix = f".{ext}"

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp.write(audio_bytes)
    tmp.close()
    return tmp.name


async def decode_input_reference(
    image_reference: ImageReference | None,
    video_reference: VideoReference | None,
    input_reference_bytes: bytes | None,
    *,
    max_video_frames: int | None = None,
    video_keep: Literal["first", "last"] = "first",
) -> Image.Image | VideoFrames | None:
    """Decode media input from multipart bytes, data URLs, or typed references."""

    provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
    if provided > 1:
        raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")

    if isinstance(input_reference_bytes, bytes):
        return _decode_media_bytes(
            input_reference_bytes,
            source="input_reference",
            max_video_frames=max_video_frames,
            video_keep=video_keep,
        )

    if isinstance(image_reference, UrlImageReference):
        return await decode_image_url(image_reference.image_url)
    elif isinstance(image_reference, FileImageReference):
        raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")

    if isinstance(video_reference, UrlVideoReference):
        return await decode_video_url(
            video_reference.video_url,
            max_frames=max_video_frames,
            keep=video_keep,
        )
    elif isinstance(video_reference, FileVideoReference):
        raise InvalidInputReferenceError("Invalid video_reference: file_id is not supported yet.")

    return None


def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
    """Normalize a torch video tensor into a numpy array of frames (F, H, W, C)."""
    video_tensor = video_tensor.detach().cpu()
    if video_tensor.dim() == 5:
        raise ValueError("Batched video tensors are not supported for single-video encoding.")
    elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
        # [C, F, H, W] -> [F, H, W, C]
        video_tensor = video_tensor.permute(1, 2, 3, 0)

    if video_tensor.is_floating_point():
        # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no
        # numpy dtype, so ``.numpy()`` below raises on it.
        video_tensor = video_tensor.float().clamp(-1, 1) * 0.5 + 0.5
    else:
        video_tensor = video_tensor.to(torch.float32) / 255.0
    video_array = video_tensor.numpy()
    return _normalize_single_video_array(video_array)


def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray:
    """Normalize a single video array into shape (F, H, W, C)."""
    if video_array.ndim == 5:
        raise ValueError("Batched video arrays are not supported for single-video encoding.")

    if video_array.ndim == 4:
        # Convert channel-first layouts to channel-last
        if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4):
            video_array = np.transpose(video_array, (1, 2, 3, 0))
        elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4):
            video_array = np.transpose(video_array, (0, 2, 3, 1))

    if np.issubdtype(video_array.dtype, np.floating):
        if video_array.min() < 0.0 or video_array.max() > 1.0:
            video_array = np.clip(video_array, -1.0, 1.0) * 0.5 + 0.5
    elif np.issubdtype(video_array.dtype, np.integer):
        video_array = video_array.astype(np.float32) / 255.0
    return video_array


def _normalize_video_array(video_array: np.ndarray) -> list[np.ndarray] | np.ndarray:
    """Normalize a numpy video array into shape (F, H, W, C).

    If a batch dimension is present, returns a list of per-video arrays.
    """
    if video_array.ndim == 5:
        return [_normalize_single_video_array(video_array[i]) for i in range(video_array.shape[0])]
    return _normalize_single_video_array(video_array)


def _normalize_frames(frames: list[Any]) -> list[np.ndarray]:
    """Normalize a list of frames into numpy arrays with values in [0,1]."""
    normalized: list[np.ndarray] = []
    for frame in frames:
        if isinstance(frame, torch.Tensor):
            frame_array = frame.detach().cpu().numpy()
        elif isinstance(frame, Image.Image):
            frame_array = np.array(frame)
        elif isinstance(frame, np.ndarray):
            frame_array = frame
        else:
            raise ValueError(f"Unsupported frame type: {type(frame)}")

        if frame_array.ndim == 3 and frame_array.shape[0] in (3, 4) and frame_array.shape[-1] not in (3, 4):
            frame_array = np.transpose(frame_array, (1, 2, 0))

        if np.issubdtype(frame_array.dtype, np.floating):
            if frame_array.min() < 0.0 or frame_array.max() > 1.0:
                frame_array = np.clip(frame_array, -1.0, 1.0) * 0.5 + 0.5
        elif np.issubdtype(frame_array.dtype, np.integer):
            frame_array = frame_array.astype(np.float32) / 255.0

        normalized.append(frame_array)
    return normalized


def _coerce_video_to_frames(video: Any) -> list[np.ndarray]:
    """Convert a video payload into a list of normalized float32 frames."""
    if isinstance(video, torch.Tensor):
        video_array = _normalize_video_tensor(video)
        return list(video_array)
    if isinstance(video, np.ndarray):
        video_array = _normalize_video_array(video)
        if isinstance(video_array, list):
            raise ValueError("Batched video arrays must be split before encoding.")
        if video_array.ndim == 4:
            return list(video_array)
        if video_array.ndim == 3:
            return [video_array]
        raise ValueError(f"Unsupported video array shape: {video_array.shape}")
    if isinstance(video, list):
        if not video:
            return []
        # If this looks like a list of frames, normalize directly.
        if all(isinstance(item, (np.ndarray, torch.Tensor, Image.Image)) for item in video):
            # If each item is itself a video (ndim==4), handle elsewhere.
            if all(hasattr(item, "ndim") and item.ndim >= 4 for item in video):
                raise ValueError("Expected a single video, got a list of video tensors/arrays.")
            return _normalize_frames(video)
        raise ValueError("Unsupported list contents for video payload.")
    raise ValueError(f"Unsupported video payload type: {type(video)}")


def _coerce_audio_to_numpy(audio: Any) -> np.ndarray:
    """Convert an audio payload into a float32 numpy array for muxing."""
    if isinstance(audio, torch.Tensor):
        arr = audio.detach().cpu().float().numpy()
    elif isinstance(audio, np.ndarray):
        arr = audio
    elif isinstance(audio, list):
        arr = np.array(audio)
    else:
        raise ValueError(f"Unsupported audio payload type: {type(audio)}")

    arr = np.squeeze(arr)
    if arr.ndim == 0:
        raise ValueError("Audio payload must contain at least one sample.")

    return arr.astype(np.float32)


def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
    """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3)."""
    frames = _coerce_video_to_frames(video)
    if not frames:
        raise ValueError("No frames found to encode.")

    frames_np = np.stack(frames, axis=0)
    if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
        frames_np = frames_np[..., :3]

    if frames_np.dtype == np.uint8:
        frames_u8 = frames_np
    else:
        frames_np = np.clip(frames_np, 0.0, 1.0)
        frames_np *= 255.0
        frames_u8 = np.round(frames_np).astype(np.uint8)

    return np.ascontiguousarray(frames_u8)


def _encode_video_bytes(
    video: Any,
    fps: int,
    audio: Any | None = None,
    audio_sample_rate: int | None = None,
    video_codec_options: dict[str, str] | None = None,
) -> bytes:
    """Encode a video payload into MP4 bytes, optionally muxing audio."""
    from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes

    audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None

    return mux_video_audio_bytes(
        _coerce_video_to_uint8_frames(video),
        audio_np,
        fps=float(fps),
        audio_sample_rate=audio_sample_rate or 24000,
        video_codec_options=video_codec_options,
    )


class FragmentedMP4VideoEncoder:
    """Normalize video chunks and append them to one fragmented MP4 stream."""

    def __init__(
        self,
        *,
        fps: int | float,
        video_codec_options: dict[str, str] | None = None,
    ) -> None:
        self._fps = float(fps)
        self._video_codec_options = video_codec_options
        self._muxer: Any | None = None

    def encode(self, video: Any) -> bytes:
        """Encode one generated video chunk and return newly emitted fMP4 bytes."""
        from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer

        frames_u8 = _coerce_video_to_uint8_frames(video)
        if self._muxer is None:
            self._muxer = FragmentedMP4Muxer(
                width=frames_u8.shape[2],
                height=frames_u8.shape[1],
                fps=self._fps,
                video_codec_options=self._video_codec_options,
            )
        return self._muxer.mux_video_frames(frames_u8)

    def close(self) -> bytes:
        """Close the underlying fMP4 muxer and return trailing bytes, if any."""
        if self._muxer is None:
            return b""
        return self._muxer.close()


StreamingVideoFormat = Literal["m4s"]


def create_streaming_video_encoder(
    *,
    output_format: StreamingVideoFormat,
    fps: int | float,
    video_codec_options: dict[str, str] | None = None,
) -> FragmentedMP4VideoEncoder:
    """Create an incremental encoder for the requested WebSocket video format."""
    if output_format == "m4s":
        return FragmentedMP4VideoEncoder(fps=fps, video_codec_options=video_codec_options)
    raise ValueError(f"Unsupported streaming video format: {output_format}")


def encode_video_base64(
    video: Any,
    fps: int,
    audio: Any | None = None,
    audio_sample_rate: int | None = None,
    video_codec_options: dict[str, str] | None = None,
) -> str:
    """Encode a video (frames/array/tensor) to base64 MP4."""
    video_bytes = _encode_video_bytes(
        video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate, video_codec_options=video_codec_options
    )
    return base64.b64encode(video_bytes).decode("utf-8")