# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""Input transform utilities for the AudioX diffusion pipeline.

Loads and normalizes the raw audio/video conditioning signals (file path / URL /
``data:`` URI / ``np.ndarray`` / ``torch.Tensor``) into the (channels, samples) and
[T, C, H, W] tensors the pipeline needs. The pipeline itself stays focused on model
forward + sampling logic.
"""

from __future__ import annotations

import base64
import os
import tempfile
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import av
import numpy as np
import soundfile
import torch
import torch.nn.functional as F
import torchaudio.functional as taF
from einops import rearrange
from torchvision.io import read_image
from vllm.multimodal.media import MediaConnector
from vllm.multimodal.media.base import MediaIO

# AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing
# "t" carry a text prompt. tv2*/v2* share the same conditioning pathways.
VIDEO_ONLY_TASKS = frozenset({"v2a", "v2m"})
TEXT_VIDEO_TASKS = frozenset({"tv2a", "tv2m"})
VIDEO_CONDITIONED_TASKS = VIDEO_ONLY_TASKS | TEXT_VIDEO_TASKS

_IMAGE_EXTS = frozenset({".jpg", ".jpeg", ".png"})


def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
    """Coerce raw prompt entries into ``{"prompt": str, ...}`` dicts (preserves extras)."""
    out: list[dict[str, Any]] = []
    for i, raw in enumerate(prompts):
        if isinstance(raw, str):
            out.append({"prompt": raw.strip()})
        elif isinstance(raw, dict):
            p = dict(raw)
            p["prompt"] = str(p.get("prompt") or "").strip()
            out.append(p)
        else:
            raise TypeError(f"AudioX prompt {i} must be str or dict, got {type(raw)!r}")
    return out


class _TempFileMediaIO(MediaIO[str]):
    """MediaIO that writes fetched bytes to a temp file, returns the path."""

    def load_bytes(self, data: bytes) -> str:
        return self._write_temp(data)

    def load_base64(self, media_type: str, data: str) -> str:
        return self._write_temp(base64.b64decode(data))

    def load_file(self, filepath: Path) -> str:
        return str(filepath)

    @staticmethod
    def _write_temp(data: bytes) -> str:
        f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
        f.write(data)
        f.close()
        return f.name


def _get_media_connector() -> MediaConnector:
    """Build a MediaConnector from the current forward context's model config."""
    from vllm_omni.diffusion.forward_context import get_forward_context

    ctx = get_forward_context()
    model_config = ctx.vllm_config.model_config if ctx.vllm_config else None
    return MediaConnector(
        allowed_local_media_path=getattr(model_config, "allowed_local_media_path", ""),
        allowed_media_domains=getattr(model_config, "allowed_media_domains", None),
    )


def materialize_media_source(source: str) -> str:
    """Return a local filesystem path for ``source``.

    Accepts a local path, a ``data:<mime>;base64,...`` URI, or an ``http(s)://`` URL.
    All inputs are routed through vLLM's ``MediaConnector`` which respects
    ``--allowed-media-domains`` and ``--allowed-local-media-path`` to prevent
    SSRF.  Bare local paths are converted to ``file://`` URIs so that the
    same path restriction applies.
    """
    if not urlparse(source).scheme:
        source = Path(source).resolve().as_uri()
    connector = _get_media_connector()
    return connector.load_from_url(source, _TempFileMediaIO())


def _load_video_path_pyav(
    path: str,
    *,
    target_fps: int,
    duration: float,
    seek_time: float,
) -> torch.Tensor:
    path = materialize_media_source(path)
    seek_time = float(seek_time)
    duration = float(duration)
    end_time = seek_time + duration if duration > 0 else None

    with av.open(path) as container:
        stream = container.streams.video[0]
        time_base = float(stream.time_base)
        src_fps = float(stream.average_rate or target_fps)

        if seek_time > 0:
            container.seek(int(seek_time / time_base), stream=stream)

        frames: list[np.ndarray] = []
        for frame in container.decode(stream):
            t = float(frame.pts) * time_base
            if t < seek_time:
                continue
            if end_time is not None and t >= end_time:
                break
            frames.append(frame.to_ndarray(format="rgb24"))

    if not frames:
        raise ValueError(f"No frames in range seek_time={seek_time!r}, duration={duration!r} for {path!r}")

    # PyAV gives [H, W, C] uint8 RGB per frame; AudioX expects [T, C, H, W].
    video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).contiguous()
    n_target = max(1, int(round(video.shape[0] * float(target_fps) / src_fps)))
    if n_target >= video.shape[0]:
        return video
    indices = torch.linspace(0, video.shape[0] - 1, n_target).round().long()
    return video[indices]


def load_video_source(
    source: Any,
    *,
    target_fps: int,
    duration: float,
    seek_time: float = 0.0,
) -> torch.Tensor:
    if isinstance(source, str):
        ext = os.path.splitext(source)[1].lower()
        if ext in _IMAGE_EXTS:
            return read_image(materialize_media_source(source)).unsqueeze(0)
        return _load_video_path_pyav(
            source,
            target_fps=target_fps,
            duration=duration,
            seek_time=seek_time,
        )

    if isinstance(source, torch.Tensor):
        return source
    if isinstance(source, np.ndarray):
        return torch.from_numpy(source)
    raise TypeError(f"Unsupported video source type: {type(source)!r}")


def normalize_video_tensor(frames: torch.Tensor, size: int = 224) -> torch.Tensor:
    if frames.dim() != 4:
        raise ValueError(f"Expected [T, C, H, W], got {tuple(frames.shape)}")

    frames = frames.float()
    if frames.max() > 1.5:
        frames = frames / 255.0

    if frames.shape[-2:] != (size, size):
        frames = F.interpolate(frames, size=(size, size), mode="bicubic", align_corners=False)
    return frames


def adjust_video_duration(frames: torch.Tensor, duration: float, target_fps: int) -> torch.Tensor:
    target_t = int(duration * target_fps)
    cur_t = frames.shape[0]

    if cur_t > target_t:
        return frames[:target_t]
    if cur_t < target_t:
        last = frames[-1:].repeat(target_t - cur_t, 1, 1, 1)
        return torch.cat([frames, last], dim=0)
    return frames


def prepare_video_reference(
    source: Any,
    *,
    duration: float,
    target_fps: int,
    seek_time: float = 0.0,
) -> torch.Tensor:
    """Decode a video clip (or single image) into the AudioX [T, 3, 224, 224] form."""
    frames = load_video_source(
        source,
        target_fps=target_fps,
        duration=duration,
        seek_time=seek_time,
    )

    if frames.dim() == 4 and frames.shape[-1] == 3:
        frames = rearrange(frames, "t h w c -> t c h w")

    frames = normalize_video_tensor(frames, size=224)
    if duration > 0:
        frames = adjust_video_duration(frames, duration, target_fps)
    return frames


def prepare_audio_reference(
    source: Any,
    *,
    model_sample_rate: int,
    seconds_start: float,
    seconds_total: float,
    device: torch.device,
) -> torch.Tensor:
    """Decode an audio source into a stereo (2, samples) tensor at the model's rate."""
    target_len = int(model_sample_rate * seconds_total)
    start = int(model_sample_rate * seconds_start)
    if isinstance(source, str):
        data, sr = soundfile.read(materialize_media_source(source), dtype="float32", always_2d=True)
        # soundfile returns channels-last (T, C); project convention is (C, T).
        wav = torch.from_numpy(data).transpose(0, 1).contiguous()
        if sr != model_sample_rate:
            wav = taF.resample(wav, sr, model_sample_rate)
    elif isinstance(source, torch.Tensor):
        wav = source.float()
    elif isinstance(source, np.ndarray):
        wav = torch.from_numpy(source).float()
    else:
        raise TypeError(f"Unsupported audio source type: {type(source)!r}")
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)
    elif wav.shape[0] > 2:
        wav = wav[:2]
    wav = wav[:, start : start + target_len]
    if wav.shape[1] < target_len:
        wav = F.pad(wav, (target_len - wav.shape[1], 0))
    return wav.to(device=device, dtype=torch.float32)


__all__ = [
    "VIDEO_ONLY_TASKS",
    "TEXT_VIDEO_TASKS",
    "VIDEO_CONDITIONED_TASKS",
    "normalize_prompts",
    "materialize_media_source",
    "load_video_source",
    "normalize_video_tensor",
    "adjust_video_duration",
    "prepare_video_reference",
    "prepare_audio_reference",
]