# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""

import io

import av
import numpy as np
import pytest

from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"

# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
_OMNI_RUNNER_PARAM = (
    MODEL,
    None,
    {
        "init_timeout": 1200,
        "tensor_parallel_size": 2,
    },
)

pytestmark = [
    pytest.mark.full_model,
    pytest.mark.diffusion,
    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]

_SKIP_ISSUE_3236 = pytest.mark.skip(
    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
)


def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
    """Validate that the MP4 contains meaningful video and audio tracks."""
    container = av.open(io.BytesIO(video_bytes))

    v_streams = [s for s in container.streams if s.type == "video"]
    assert len(v_streams) >= 1, "No video stream found in MP4"

    a_streams = [s for s in container.streams if s.type == "audio"]
    assert len(a_streams) >= 1, "No audio stream found in MP4"

    v_stream = v_streams[0]
    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"

    frame_count = 0
    for frame in container.decode(video=0):
        frame_count += 1
        if frame_count >= min_frames:
            break
    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"

    container.close()


@_SKIP_ISSUE_3236
@hardware_test(res={"cuda": "H100"}, num_cards=2)
def test_magi_human_e2e(omni_runner: OmniRunner):
    """End-to-end test for MagiHuman generating video and audio."""
    prompt = (
        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
        "stationary in front of a softly lit, blurred warm-toned interior. Her "
        "overall disposition is warm, composed, and gently confident. The camera "
        "holds a static medium close-up, framing her from the shoulders up, "
        "with shallow depth of field keeping her face in sharp focus. Soft "
        "directional key light falls from the upper left, casting a gentle "
        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
        "the levator labii superiors relaxing as her lips part. She speaks in "
        "clear, warm, unhurried American English: "
        "\"The most beautiful things in life aren't things at all — "
        "they're moments, feelings, and the people who make you feel truly alive.\" "
        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
        "shapes each vowel with precision. A faint, genuine smile engages the "
        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
        "Dialogue:\n"
        "<Young blonde woman, American English>: "
        "\"The most beautiful things in life aren't things at all — "
        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
        "Background Sound:\n"
        "<Soft, warm indoor ambience with a faint distant piano melody>"
    )

    sampling_params = OmniDiffusionSamplingParams(
        height=256,
        width=448,
        num_inference_steps=8,
        seed=52,
        extra_args={
            "seconds": 5,
            "sr_height": 1080,
            "sr_width": 1920,
            "sr_num_inference_steps": 5,
        },
    )

    omni = omni_runner.omni
    outputs = list(
        omni.generate(
            prompts=[prompt],
            sampling_params_list=[sampling_params],
        )
    )

    assert len(outputs) > 0, "No outputs returned"
    first = outputs[0]

    assert hasattr(first, "images") and first.images, "No video frames in output"
    video_frames = first.images[0]
    assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
    assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"

    mm = first.multimodal_output
    assert mm, "multimodal_output is empty or missing"

    audio_waveform = mm.get("audio")
    assert audio_waveform is not None, "No audio waveform in multimodal_output"

    audio_sample_rate = mm.get("audio_sample_rate")
    assert audio_sample_rate is not None, (
        "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
    )
    assert isinstance(audio_sample_rate, (int, float)), (
        f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
    )
    assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"

    fps = mm.get("fps")
    assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
    assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
    assert int(fps) > 0, f"fps must be positive, got {fps}"

    video_bytes = mux_video_audio_bytes(
        video_frames,
        audio_waveform,
        fps=float(fps),
        audio_sample_rate=int(audio_sample_rate),
    )
    assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
    assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"

    _validate_mp4(video_bytes)