#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Gradio demo for `/v1/realtime/video`.

Usage:
    python gradio_demo.py --host 127.0.0.1 --port 7860

Start a vLLM-Omni server with streaming output enabled first, for example:
    vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any, TypedDict

try:
    import gradio as gr
except ImportError:
    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None

DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement."
DEFAULT_TRANSITION_CHUNKS = 3
PROMPT_UPDATE_HEADERS = ["At (s)", "Prompt", "Transition Chunks"]
PROMPT_UPDATE_VALUES = [
    [
        4,
        (
            "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. "
            "The water is murky, swirling violently, carrying debris and marine life into the vortex. "
            "The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. "
            "The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. "
            "Close-up shot emphasizing the turbulent motion and destruction."
        ),
        3,
    ],
    [
        11,
        (
            "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. "
            "Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. "
            "The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. "
            "Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. "
            "Close-up shot emphasizing the turbulent motion and destruction."
        ),
        3,
    ],
]
VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html")
VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js")


class ScheduledPromptUpdate(TypedDict):
    at: float
    prompt: str
    transition_chunks: int


def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
    if value is None:
        return
    if isinstance(value, str):
        payload[key] = value.strip()
    else:
        payload[key] = value


def _optional_int(value: int | float | None) -> int | None:
    return None if value is None else int(value)


def _optional_float(value: int | float | None) -> float | None:
    return None if value is None else float(value)


def _normalize_prompt_update_row(row: list[Any] | tuple[Any, ...]) -> tuple[Any, Any, Any]:
    values = []
    for v in row[: len(PROMPT_UPDATE_HEADERS)]:
        if isinstance(v, str):
            values.append(v.strip() or None)
        elif isinstance(v, float) and math.isnan(v):
            values.append(None)
        else:
            values.append(v)
    while len(values) < len(PROMPT_UPDATE_HEADERS):
        values.append(None)
    return values[0], values[1], values[2]  # pyright: ignore[reportIndexIssue]


def _prompt_update_row_is_empty(row: list[Any] | tuple[Any, ...]) -> bool:
    at, prompt, transition = _normalize_prompt_update_row(row)
    return at is None and prompt is None and transition is None


def _parse_prompt_update_row(row: list[Any] | tuple[Any, ...], *, index: int) -> ScheduledPromptUpdate:
    at_val, prompt_val, transition_val = _normalize_prompt_update_row(row)
    row_label = f"Prompt update row {index + 1}"

    if at_val is None or prompt_val is None:
        raise gr.Error(f"{row_label} requires both At (s) and Prompt.")

    try:
        at = float(at_val)
    except (TypeError, ValueError) as exc:
        raise gr.Error(f"{row_label}: At (s) must be a number.") from exc
    if at < 0:
        raise gr.Error(f"{row_label}: At (s) must be >= 0.")

    prompt = str(prompt_val).strip()
    if not prompt:
        raise gr.Error(f"{row_label}: Prompt must be a non-empty string.")

    transition_chunks = DEFAULT_TRANSITION_CHUNKS
    if transition_val is not None:
        try:
            transition_chunks = int(transition_val)
        except (TypeError, ValueError) as exc:
            raise gr.Error(f"{row_label}: Transition Chunks must be an integer.") from exc
        if transition_chunks < 0:
            raise gr.Error(f"{row_label}: Transition Chunks must be >= 0.")

    return {
        "at": at,
        "prompt": prompt,
        "transition_chunks": transition_chunks,
    }


def _parse_prompt_updates_dataframe(rows: list[list[Any]] | None) -> list[ScheduledPromptUpdate]:
    if not rows:
        return []

    # Single row: empty or complete.
    if len(rows) == 1:
        row = rows[0]
        if _prompt_update_row_is_empty(row):
            return []
        return [_parse_prompt_update_row(row, index=0)]

    updates: list[ScheduledPromptUpdate] = []
    for index, row in enumerate(rows):
        if _prompt_update_row_is_empty(row):
            raise gr.Error(f"Prompt update row {index + 1} cannot be empty when multiple rows are present.")
        updates.append(_parse_prompt_update_row(row, index=index))
    return sorted(updates, key=lambda update: update["at"])


def _build_session_start(
    *,
    model: str,
    prompt: str,
    negative_prompt: str | None,
    width: int | float | None,
    height: int | float | None,
    fps: int | float | None,
    num_frames: int | float | None,
    num_inference_steps: int | float | None,
    guidance_scale: int | float | None,
    guidance_scale_2: int | float | None,
    boundary_ratio: int | float | None,
    flow_shift: int | float | None,
    true_cfg_scale: int | float | None,
    seed: int | float | None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "type": "session.start",
        "model": model,
        "prompt": prompt,
        "format": "m4s",
    }

    _maybe_set(payload, "negative_prompt", negative_prompt)
    _maybe_set(payload, "width", _optional_int(width))
    _maybe_set(payload, "height", _optional_int(height))
    _maybe_set(payload, "fps", _optional_int(fps))
    _maybe_set(payload, "num_frames", _optional_int(num_frames))
    _maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps))
    _maybe_set(payload, "guidance_scale", _optional_float(guidance_scale))
    _maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2))
    _maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio))
    _maybe_set(payload, "flow_shift", _optional_float(flow_shift))
    _maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale))
    _maybe_set(payload, "seed", _optional_int(seed))

    if "Helios" in model:
        payload["extra_params"] = {
            "is_enable_stage2": True,
            "pyramid_num_stages": 3,
            "pyramid_num_inference_steps_list": [1, 1, 1],
            "is_amplify_first_chunk": True,
        }

    return payload


def _set_optional_field_enabled(enabled: bool) -> Any:
    return gr.update(interactive=enabled)


def build_browser_stream_config(
    host: str,
    port: int | float,
    model: str,
    prompt: str,
    negative_prompt: str,
    include_negative_prompt: bool,
    width: int | float | None,
    height: int | float | None,
    fps: int | float,
    num_frames: int | float | None,
    num_inference_steps: int | float | None,
    include_guidance_scale: bool,
    guidance_scale: int | float | None,
    include_guidance_scale_2: bool,
    guidance_scale_2: int | float | None,
    include_seed: bool,
    seed: int | float | None,
    include_boundary_ratio: bool,
    boundary_ratio: int | float | None,
    include_flow_shift: bool,
    flow_shift: int | float | None,
    include_true_cfg_scale: bool,
    true_cfg_scale: int | float | None,
    prompt_updates_df: list[list[Any]] | None,
) -> tuple[str, Any]:
    """Build the browser-side WebSocket request config for the HTML/MSE player."""
    if not model.strip():
        raise gr.Error("Model is required.")
    if not prompt.strip():
        raise gr.Error("Prompt is required.")

    payload = _build_session_start(
        model=model.strip(),
        prompt=prompt.strip(),
        negative_prompt=negative_prompt if include_negative_prompt else None,
        width=width,
        height=height,
        fps=fps,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale if include_guidance_scale else None,
        guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None,
        boundary_ratio=boundary_ratio if include_boundary_ratio else None,
        flow_shift=flow_shift if include_flow_shift else None,
        true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None,
        seed=seed if include_seed else None,
    )
    prompt_updates = _parse_prompt_updates_dataframe(prompt_updates_df)
    config = {
        "url": f"ws://{host}:{int(port)}/v1/realtime/video",
        "payload": payload,
    }
    if prompt_updates:
        config["prompt_updates"] = prompt_updates
    return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False)


def create_demo() -> gr.Blocks:
    with gr.Blocks(
        title="Streaming Video Generation",
        css="textarea:disabled, input:disabled { color: var(--body-text-color-subdued); -webkit-text-fill-color: var(--body-text-color-subdued); background-color: var(--background-fill-secondary); }",
    ) as demo:
        gr.Markdown("# Streaming Video Generation")
        gr.Markdown(
            "Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks "
            "directly to a Media Source Extensions video player."
        )

        with gr.Row():
            with gr.Column(scale=1):
                host = gr.Textbox(label="Server Host", value="127.0.0.1")
                port = gr.Number(label="Server Port", value=8000, precision=0)
                model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled")
                prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)

                with gr.Row():
                    width = gr.Number(label="Width", value=640, precision=0)
                    height = gr.Number(label="Height", value=384, precision=0)

                with gr.Row():
                    fps = gr.Number(label="FPS", value=16, precision=0)
                    num_frames = gr.Number(label="Num Frames", value=429, precision=0)
                    num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0)

                with gr.Accordion("Advanced Sampling", open=False):
                    with gr.Row():
                        with gr.Column():
                            negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=1, interactive=False)
                            include_negative_prompt = gr.Checkbox(label="Send", value=False)
                    with gr.Row():
                        with gr.Column(scale=0, min_width=150):
                            guidance_scale = gr.Number(
                                label="Guidance Scale",
                                value=1.0,
                                interactive=True,
                            )
                            include_guidance_scale = gr.Checkbox(label="Send", value=True)
                        with gr.Column(scale=0, min_width=150):
                            guidance_scale_2 = gr.Number(
                                label="Guidance Scale 2",
                                value=None,
                                interactive=False,
                            )
                            include_guidance_scale_2 = gr.Checkbox(label="Send", value=False)
                        with gr.Column(scale=0, min_width=150):
                            true_cfg_scale = gr.Number(
                                label="True CFG Scale",
                                value=None,
                                interactive=False,
                            )
                            include_true_cfg_scale = gr.Checkbox(label="Send", value=False)
                        with gr.Column(scale=0, min_width=150):
                            boundary_ratio = gr.Number(
                                label="Boundary Ratio",
                                value=None,
                                interactive=False,
                            )
                            include_boundary_ratio = gr.Checkbox(label="Send", value=False)
                        with gr.Column(scale=0, min_width=150):
                            flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False)
                            include_flow_shift = gr.Checkbox(label="Send", value=False)
                        with gr.Column(scale=0, min_width=150):
                            seed = gr.Number(label="Seed", value=42, precision=0, interactive=True)
                            include_seed = gr.Checkbox(label="Send", value=True)

                with gr.Accordion("Prompt Updates", open=True):
                    gr.Markdown(
                        "Schedule midway text `session.interaction` messages after `video.start`. "
                        "With one row, leave it fully empty or fill both **At (s)** and **Prompt**. "
                        "With multiple rows, every row must be complete."
                    )
                    prompt_updates_df = gr.Dataframe(
                        headers=PROMPT_UPDATE_HEADERS,
                        datatype=["number", "str", "number"],
                        type="array",
                        row_count=1,
                        row_limits=(1, None),
                        column_count=3,
                        value=PROMPT_UPDATE_VALUES,
                        interactive=True,
                    )

                with gr.Row():
                    start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start")
                    hidden_stream_config = gr.Textbox(value="", visible=False)

            with gr.Column(scale=1):
                gr.HTML(
                    VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"),
                    label="Streaming Preview",
                    js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"),
                )

        inputs = [
            host,
            port,
            model,
            prompt,
            negative_prompt,
            include_negative_prompt,
            width,
            height,
            fps,
            num_frames,
            num_inference_steps,
            include_guidance_scale,
            guidance_scale,
            include_guidance_scale_2,
            guidance_scale_2,
            include_seed,
            seed,
            include_boundary_ratio,
            boundary_ratio,
            include_flow_shift,
            flow_shift,
            include_true_cfg_scale,
            true_cfg_scale,
            prompt_updates_df,
        ]
        for toggle, field in [
            (include_negative_prompt, negative_prompt),
            (include_guidance_scale, guidance_scale),
            (include_guidance_scale_2, guidance_scale_2),
            (include_seed, seed),
            (include_boundary_ratio, boundary_ratio),
            (include_flow_shift, flow_shift),
            (include_true_cfg_scale, true_cfg_scale),
        ]:
            toggle.change(
                fn=_set_optional_field_enabled,
                inputs=toggle,
                outputs=field,
            )
        start_button.click(
            fn=build_browser_stream_config,
            inputs=inputs,
            outputs=[hidden_stream_config, start_button],
        ).then(
            fn=lambda config: config,
            inputs=hidden_stream_config,
            outputs=hidden_stream_config,
            js="""
            (config) => {
              if (window.vllmStreamingVideoStart) {
                window.vllmStreamingVideoStart(config);
              }
              return config;
            }
            """,
        )

    return demo


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo")
    parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch")
    parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch")
    parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    demo = create_demo()
    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share)


if __name__ == "__main__":
    main()