import argparse
import os
import random
import signal
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Optional

import gradio as gr
import numpy as np
import soundfile as sf
import torch
from PIL import Image
from vllm.assets.video import video_get_metadata, video_to_ndarrays
from vllm.sampling_params import SamplingParams

from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM

# Import utils from offline inference example
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
try:
    from utils import make_omni_prompt
except ImportError:
    # Fallback if utils doesn't exist - we'll build prompts directly
    make_omni_prompt = None

SEED = 42
ASYNC_INIT_TIMEOUT = 600

SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
    "Qwen/Qwen2.5-Omni-7B": {
        "sampling_params": {
            "thinker": {
                "temperature": 0.0,
                "top_p": 1.0,
                "top_k": -1,
                "max_tokens": 2048,
                "detokenize": True,
                "repetition_penalty": 1.1,
            },
            "talker": {
                "temperature": 0.0,
                "top_p": 1.0,
                "top_k": -1,
                "max_tokens": 2048,
                "detokenize": True,
                "repetition_penalty": 1.1,
                "stop_token_ids": [8294],
            },
            "code2wav": {
                "temperature": 0.0,
                "top_p": 1.0,
                "top_k": -1,
                "max_tokens": 2048,
                "detokenize": True,
                "repetition_penalty": 1.1,
            },
        },
    },
}
# Ensure deterministic behavior across runs.
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
os.environ["PYTHONHASHSEED"] = str(SEED)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"


def parse_args():
    parser = argparse.ArgumentParser(description="Gradio demo for Qwen2.5-Omni online inference.")
    parser.add_argument(
        "--model",
        default="Qwen/Qwen2.5-Omni-7B",
        help="Path to local model checkpoint.",
    )
    parser.add_argument(
        "--ip",
        default="127.0.0.1",
        help="Host/IP for gradio `launch`.",
    )
    parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.")
    parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
    parser.add_argument(
        "--stage-configs-path",
        type=str,
        default=None,
        help="Path to custom stage configs YAML file (optional).",
    )
    return parser.parse_args()


def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
    """Build SamplingParams objects by reusing the dict definitions."""
    return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]


def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]:
    """Build sampling params as dict for HTTP API mode."""
    model_conf = SUPPORTED_MODELS.get(model_key)
    if model_conf is None:
        raise ValueError(f"Unsupported model '{model_key}'")

    sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"]
    sampling_params: list[dict] = []
    for stage_name, template in sampling_templates.items():
        params = dict(template)
        params["seed"] = seed
        sampling_params.append(params)
    return sampling_params


def create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace:
    # The prompt builder expects a minimal namespace with these attributes.
    return SimpleNamespace(
        model=base_args.model,
        prompt_type="text",
        tokenize=True,
        use_torchvision=True,
        legacy_omni_video=False,
    )


def process_audio_file(
    audio_file: Optional[Any],
) -> Optional[tuple[np.ndarray, int]]:
    """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
    if audio_file is None:
        return None

    sample_rate: Optional[int] = None
    audio_np: Optional[np.ndarray] = None

    def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]:
        if not path_str:
            return None
        path = Path(path_str)
        if not path.exists():
            return None
        data, sr = sf.read(path)
        if data.ndim > 1:
            data = data[:, 0]
        return data.astype(np.float32), int(sr)

    if isinstance(audio_file, tuple):
        if len(audio_file) == 2:
            first, second = audio_file
            # Case 1: (sample_rate, np.ndarray)
            if isinstance(first, (int, float)) and isinstance(second, np.ndarray):
                sample_rate = int(first)
                audio_np = second
            # Case 2: (filepath, (sample_rate, np.ndarray or list))
            elif isinstance(first, str):
                if isinstance(second, tuple) and len(second) == 2:
                    sr_candidate, data_candidate = second
                    if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
                        sample_rate = int(sr_candidate)
                        audio_np = data_candidate
                if audio_np is None:
                    loaded = _load_from_path(first)
                    if loaded is not None:
                        audio_np, sample_rate = loaded
            # Case 3: (None, (sample_rate, np.ndarray))
            elif first is None and isinstance(second, tuple) and len(second) == 2:
                sr_candidate, data_candidate = second
                if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
                    sample_rate = int(sr_candidate)
                    audio_np = data_candidate
        elif len(audio_file) == 1 and isinstance(audio_file[0], str):
            loaded = _load_from_path(audio_file[0])
            if loaded is not None:
                audio_np, sample_rate = loaded
    elif isinstance(audio_file, str):
        loaded = _load_from_path(audio_file)
        if loaded is not None:
            audio_np, sample_rate = loaded

    if audio_np is None or sample_rate is None:
        return None

    if audio_np.ndim > 1:
        audio_np = audio_np[:, 0]

    return audio_np.astype(np.float32), sample_rate


def process_image_file(image_file: Optional[Image.Image]) -> Optional[Image.Image]:
    """Process image file from Gradio input.

    Returns:
        PIL Image in RGB mode or None if no image provided.
    """
    if image_file is None:
        return None
    # Convert to RGB if needed
    if image_file.mode != "RGB":
        image_file = image_file.convert("RGB")
    return image_file


def process_video_file(
    video_file: Optional[str],
    enable_audio_in_video: bool = False,
    max_frames: int = 32,
) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]:
    """Process video file and optionally extract audio track."""
    if video_file is None:
        return None

    video_path = Path(video_file)
    if not video_path.exists():
        print(f"Video file not found: {video_path}")
        return None

    try:
        frames = video_to_ndarrays(str(video_path), num_frames=max_frames)
        metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0])
    except Exception as exc:  # pylint: disable=broad-except
        print(f"Failed to decode video {video_path}: {exc}")
        return None

    audio_tuple: Optional[tuple[np.ndarray, int]] = None
    if enable_audio_in_video:
        try:
            import librosa  # type: ignore import

            audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000)
            audio_tuple = (audio_signal.astype(np.float32), sampling_rate)
        except Exception as exc:  # pylint: disable=broad-except
            print(f"Failed to extract audio from video {video_path}: {exc}")

    return frames, metadata, audio_tuple


async def run_inference_async_omni_llm(
    omni_llm: AsyncOmniLLM,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    user_prompt: str,
    audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None,
    image_file: Optional[Image.Image] = None,
    video_file: Optional[str] = None,
    use_audio_in_video: bool = False,
):
    """Run inference using AsyncOmniLLM directly with multimodal support."""
    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
        return "Please provide at least a text prompt or multimodal input.", None

    try:
        # Build prompt with multimodal data
        prompt_args = SimpleNamespace(**prompt_args_template.__dict__)

        # Process multimodal inputs
        multi_modal_data = {}
        mm_processor_kwargs = {}

        # Process audio
        audio_data = process_audio_file(audio_file)
        if audio_data is not None:
            multi_modal_data["audio"] = audio_data

        # Process image
        image_data = process_image_file(image_file)
        if image_data is not None:
            multi_modal_data["image"] = image_data

        # Process video
        video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video)
        if video_payload is not None:
            video_frames, video_metadata, extracted_audio = video_payload
            video_entry: Any
            if video_metadata:
                video_entry = (video_frames, video_metadata)
            else:
                video_entry = video_frames
            multi_modal_data["video"] = video_entry
            if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
                multi_modal_data["audio"] = extracted_audio
                mm_processor_kwargs["use_audio_in_video"] = True

        # Build the prompt input
        if make_omni_prompt is not None:
            omni_prompt = make_omni_prompt(prompt_args, user_prompt)
            # Add multimodal data if present
            if multi_modal_data:
                if isinstance(omni_prompt, dict):
                    omni_prompt["multi_modal_data"] = multi_modal_data
                    if mm_processor_kwargs:
                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
                else:
                    # If make_omni_prompt returns a string, we need to wrap it
                    omni_prompt = {
                        "prompt": omni_prompt,
                        "multi_modal_data": multi_modal_data,
                    }
                    if mm_processor_kwargs:
                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
        else:
            # Fallback: build prompt directly
            default_system = (
                "You are Qwen, a virtual human developed by the Qwen Team, "
                "Alibaba Group, capable of perceiving auditory and visual inputs, "
                "as well as generating text and speech."
            )
            prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
            if audio_data:
                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
            if image_data:
                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
            if video_payload is not None:
                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
            if user_prompt.strip():
                prompt += f"{user_prompt}"
            prompt += "<|im_end|>\n<|im_start|>assistant\n"

            omni_prompt = {
                "prompt": prompt,
                "multi_modal_data": multi_modal_data,
            }
            if mm_processor_kwargs:
                omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs

        request_id = "0"
        text_outputs: list[str] = []
        audio_output = None

        async for stage_outputs in omni_llm.generate(
            prompt=omni_prompt,
            request_id=request_id,
            sampling_params_list=sampling_params,
        ):
            # stage_outputs.request_output is a RequestOutput object, not a list
            request_output = stage_outputs.request_output
            if stage_outputs.final_output_type == "text":
                if request_output.outputs:
                    for output in request_output.outputs:
                        if output.text:
                            text_outputs.append(output.text)
            elif stage_outputs.final_output_type == "audio":
                # multimodal_output is on the RequestOutput object
                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
                if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
                    if "audio" in request_output.multimodal_output:
                        audio_tensor = request_output.multimodal_output["audio"]
                        # Ensure audio is 1D (flatten if needed)
                        if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
                            audio_tensor = audio_tensor.flatten()
                        audio_np = audio_tensor.detach().cpu().numpy()
                        audio_output = (
                            24000,  # sampling rate in Hz
                            audio_np,
                        )

        text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
        return text_response, audio_output
    except Exception as exc:  # pylint: disable=broad-except
        return f"Inference failed: {exc}", None


def build_interface(
    omni_llm: AsyncOmniLLM,
    sampling_params: list[SamplingParams],
    prompt_args_template: SimpleNamespace,
    model: str,
):
    """Build Gradio interface for AsyncOmniLLM mode."""

    async def run_inference(
        user_prompt: str,
        audio_file: Optional[tuple[str, tuple[int, np.ndarray]]],
        image_file: Optional[Image.Image],
        video_file: Optional[str],
        use_audio_in_video: bool,
    ):
        return await run_inference_async_omni_llm(
            omni_llm,
            sampling_params,
            prompt_args_template,
            user_prompt,
            audio_file,
            image_file,
            video_file,
            use_audio_in_video,
        )

    css = """
    .media-input-container {
        display: flex;
        gap: 10px;
    }
    .media-input-container > div {
        flex: 1;
    }
    .media-input-container .image-input,
    .media-input-container .audio-input {
        height: 300px;
    }
    .media-input-container .video-column {
        height: 300px;
        display: flex;
        flex-direction: column;
    }
    .media-input-container .video-input {
        flex: 1;
        min-height: 0;
    }
    #generate-btn button {
        width: 100%;
    }
    """

    with gr.Blocks(css=css) as demo:
        gr.Markdown("# vLLM-Omni Online Serving Demo")
        gr.Markdown(f"**Model:** {model} \n\n")

        with gr.Column():
            with gr.Row():
                input_box = gr.Textbox(
                    label="Text Prompt",
                    placeholder="For example: Describe what happens in the media inputs.",
                    lines=4,
                    scale=1,
                )
            with gr.Row(elem_classes="media-input-container"):
                image_input = gr.Image(
                    label="Image Input (optional)",
                    type="pil",
                    sources=["upload"],
                    scale=1,
                    elem_classes="image-input",
                )
                with gr.Column(scale=1, elem_classes="video-column"):
                    video_input = gr.Video(
                        label="Video Input (optional)",
                        sources=["upload"],
                        elem_classes="video-input",
                    )
                    use_audio_in_video_checkbox = gr.Checkbox(
                        label="Use audio from video",
                        value=False,
                        info="Extract the video's audio track when provided.",
                    )
                audio_input = gr.Audio(
                    label="Audio Input (optional)",
                    type="numpy",
                    sources=["upload", "microphone"],
                    scale=1,
                    elem_classes="audio-input",
                )

        with gr.Row():
            generate_btn = gr.Button(
                "Generate",
                variant="primary",
                size="lg",
                elem_id="generate-btn",
            )

        with gr.Row():
            text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
            audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)

        generate_btn.click(
            fn=run_inference,
            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
            outputs=[text_output, audio_output],
        )
        demo.queue()
    return demo


def main():
    args = parse_args()
    omni_llm = None

    model_name = "/".join(args.model.split("/")[-2:])
    assert model_name in SUPPORTED_MODELS, (
        f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
    )

    # Register signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        print("\nReceived interrupt signal, shutting down...")
        if omni_llm is not None:
            try:
                omni_llm.shutdown()
            except Exception as e:
                print(f"Error during shutdown: {e}")
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    print(f"Initializing AsyncOmniLLM with model: {args.model}")
    if args.stage_configs_path:
        print(f"Using custom stage configs: {args.stage_configs_path}")

    sampling_params = build_sampling_params(SEED, model_name)
    omni_llm = AsyncOmniLLM(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
        init_timeout=ASYNC_INIT_TIMEOUT,
    )
    print("✓ AsyncOmniLLM initialized successfully")
    prompt_args_template = create_prompt_args(args)

    demo = build_interface(
        omni_llm,
        sampling_params,
        prompt_args_template,
        args.model,
    )
    try:
        demo.launch(
            server_name=args.ip,
            server_port=args.port,
            share=args.share,
        )
    except KeyboardInterrupt:
        print("\nShutting down...")
    finally:
        # Cleanup
        if omni_llm is not None:
            try:
                omni_llm.shutdown()
            except Exception as e:
                print(f"Error during cleanup: {e}")


if __name__ == "__main__":
    main()