# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Offline end-to-end inference for MiniCPM-o 4.5 (thinker + talker/Token2Wav).

MiniCPM-o uses placeholders ``(<image>./</image>)``, ``(<audio>./</audio>)``,
and ``(<video>./</video>)``. Speech output requires a ``<|tts_bos|>`` suffix on
the assistant prefix (equivalent to online ``chat_template_kwargs.use_tts_template``).
"""

import os
import time
from typing import NamedTuple

import numpy as np
import soundfile as sf
import vllm
from PIL import Image
from vllm import SamplingParams
from vllm.assets.audio import AudioAsset
from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio

from vllm_omni.entrypoints.omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser

SEED = 42

default_system = (
    "You are MiniCPM-o, a helpful multimodal assistant that can "
    "understand images, audio and video, and respond in text and speech."
)


class QueryResult(NamedTuple):
    inputs: dict
    limit_mm_per_prompt: dict[str, int]


def _assistant_prefix(use_tts: bool) -> str:
    # Matches HF chat_template: assistant header, optional empty <think>, then TTS bos.
    prefix = "<|im_start|>assistant\n"
    if use_tts:
        prefix += "<|tts_bos|>"
    return prefix


def _build_prompt(user_body: str, use_tts: bool) -> str:
    return (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
        f"<|im_start|>user\n{user_body}<|im_end|>\n"
        f"{_assistant_prefix(use_tts)}"
    )


def get_text_query(question: str | None = None, use_tts: bool = True) -> QueryResult:
    if question is None:
        question = "Say hello, then introduce vLLM-Omni in one sentence."
    return QueryResult(
        inputs={"prompt": _build_prompt(question, use_tts=use_tts)},
        limit_mm_per_prompt={},
    )


def get_image_query(
    question: str | None = None,
    image_path: str | None = None,
    use_tts: bool = True,
) -> QueryResult:
    if question is None:
        question = "What is the content of this image?"
    user_body = f"(<image>./</image>)\n{question}"

    if image_path:
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image file not found: {image_path}")
        image_data = convert_image_mode(Image.open(image_path), "RGB")
    else:
        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")

    return QueryResult(
        inputs={
            "prompt": _build_prompt(user_body, use_tts=use_tts),
            "multi_modal_data": {"image": image_data},
        },
        limit_mm_per_prompt={"image": 1},
    )


def get_audio_query(
    question: str | None = None,
    audio_path: str | None = None,
    sampling_rate: int = 16000,
    use_tts: bool = True,
) -> QueryResult:
    if question is None:
        question = "What is the content of this audio?"
    user_body = f"(<audio>./</audio>)\n{question}"

    if audio_path:
        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file not found: {audio_path}")
        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
        audio_data = (audio_signal.astype(np.float32), sr)
    else:
        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate

    return QueryResult(
        inputs={
            "prompt": _build_prompt(user_body, use_tts=use_tts),
            "multi_modal_data": {"audio": audio_data},
        },
        limit_mm_per_prompt={"audio": 1},
    )


def get_video_query(
    question: str | None = None,
    video_path: str | None = None,
    num_frames: int = 16,
    use_tts: bool = True,
) -> QueryResult:
    if question is None:
        question = "Why is this video funny?"
    user_body = f"(<video>./</video>)\n{question}"

    if video_path:
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video file not found: {video_path}")
        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
    else:
        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays

    return QueryResult(
        inputs={
            "prompt": _build_prompt(user_body, use_tts=use_tts),
            "multi_modal_data": {"video": video_frames},
        },
        limit_mm_per_prompt={"video": 1},
    )


def get_mixed_modalities_query(
    video_path: str | None = None,
    image_path: str | None = None,
    audio_path: str | None = None,
    num_frames: int = 16,
    sampling_rate: int = 16000,
    use_tts: bool = True,
) -> QueryResult:
    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
    user_body = f"(<audio>./</audio>)\n(<image>./</image>)\n(<video>./</video>)\n{question}"

    if video_path:
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video file not found: {video_path}")
        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
    else:
        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays

    if image_path:
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image file not found: {image_path}")
        image_data = convert_image_mode(Image.open(image_path), "RGB")
    else:
        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")

    if audio_path:
        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file not found: {audio_path}")
        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
        audio_data = (audio_signal.astype(np.float32), sr)
    else:
        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate

    return QueryResult(
        inputs={
            "prompt": _build_prompt(user_body, use_tts=use_tts),
            "multi_modal_data": {
                "audio": audio_data,
                "image": image_data,
                "video": video_frames,
            },
        },
        limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
    )


def get_multi_audios_query(use_tts: bool = True) -> QueryResult:
    question = "Are these two audio clips the same?"
    user_body = f"(<audio>./</audio>)\n(<audio>./</audio>)\n{question}"
    return QueryResult(
        inputs={
            "prompt": _build_prompt(user_body, use_tts=use_tts),
            "multi_modal_data": {
                "audio": [
                    AudioAsset("winning_call").audio_and_sample_rate,
                    AudioAsset("mary_had_lamb").audio_and_sample_rate,
                ],
            },
        },
        limit_mm_per_prompt={"audio": 2},
    )


query_map = {
    "text": get_text_query,
    "use_audio": get_audio_query,
    "use_image": get_image_query,
    "use_video": get_video_query,
    "use_multi_audios": get_multi_audios_query,
    "use_mixed_modalities": get_mixed_modalities_query,
}


def _wants_tts(modalities: str | None) -> bool:
    """TTS bos is needed whenever audio is among the requested output modalities."""
    if modalities is None:
        return True
    parts = {p.strip() for p in modalities.split(",") if p.strip()}
    return "audio" in parts


def main(args):
    model_name = args.model
    print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)

    video_path = getattr(args, "video_path", None)
    image_path = getattr(args, "image_path", None)
    audio_path = getattr(args, "audio_path", None)
    use_tts = _wants_tts(args.modalities)

    query_func = query_map[args.query_type]
    if args.query_type == "use_video":
        query_result = query_func(
            video_path=video_path,
            num_frames=getattr(args, "num_frames", 16),
            use_tts=use_tts,
        )
    elif args.query_type == "use_image":
        query_result = query_func(image_path=image_path, use_tts=use_tts)
    elif args.query_type == "use_audio":
        query_result = query_func(
            audio_path=audio_path,
            sampling_rate=getattr(args, "sampling_rate", 16000),
            use_tts=use_tts,
        )
    elif args.query_type == "use_mixed_modalities":
        query_result = query_func(
            video_path=video_path,
            image_path=image_path,
            audio_path=audio_path,
            num_frames=getattr(args, "num_frames", 16),
            sampling_rate=getattr(args, "sampling_rate", 16000),
            use_tts=use_tts,
        )
    elif args.query_type == "use_multi_audios":
        query_result = query_func(use_tts=use_tts)
    else:
        query_result = query_func(use_tts=use_tts)

    omni_kwargs = vars(args).copy()
    # Drop example-only CLI keys that are not Omni / engine args.
    for key in (
        "query_type",
        "output_wav",
        "output_dir",
        "num_prompts",
        "txt_prompts",
        "video_path",
        "image_path",
        "audio_path",
        "num_frames",
        "sampling_rate",
        "modalities",
        "py_generator",
        "enable_profiler",
        "profiler_stages",
        "log_dir",
    ):
        omni_kwargs.pop(key, None)
    omni_kwargs["model"] = model_name
    omni_kwargs["trust_remote_code"] = True
    omni = Omni(**omni_kwargs)

    # Stage 0 (thinker): multimodal understanding → text (+ TTS span when enabled).
    thinker_sampling_params = SamplingParams(
        temperature=0.0,
        top_p=1.0,
        top_k=-1,
        max_tokens=2048,
        seed=SEED,
        detokenize=True,
        repetition_penalty=1.1,
    )
    # Stage 1 (talker + Token2Wav): max_tokens=1 satisfies the scheduler;
    # waveform is produced in-process by Token2Wav.
    talker_sampling_params = SamplingParams(
        temperature=0.0,
        top_p=1.0,
        top_k=-1,
        max_tokens=1,
        seed=SEED,
        detokenize=False,
    )
    sampling_params_list = [thinker_sampling_params, talker_sampling_params][: omni.num_stages]

    if args.txt_prompts is None:
        prompts = [query_result.inputs for _ in range(args.num_prompts)]
    else:
        assert args.query_type == "text", "txt-prompts is only supported for text query type"
        with open(args.txt_prompts, encoding="utf-8") as f:
            lines = [ln.strip() for ln in f.readlines()]
            prompts = [get_text_query(ln, use_tts=use_tts).inputs for ln in lines if ln != ""]
            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")

    if args.modalities is not None:
        output_modalities = args.modalities.split(",")
        for prompt in prompts:
            prompt["modalities"] = output_modalities

    profiler_enabled = args.enable_profiler
    if profiler_enabled:
        omni.start_profile(stages=args.profiler_stages)

    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
    os.makedirs(output_dir, exist_ok=True)

    total_requests = len(prompts)
    processed_count = 0
    print(f"query type: {args.query_type}")
    print(f"use_tts (tts_bos): {use_tts}")

    for stage_outputs in omni_generator:
        output = stage_outputs
        if stage_outputs.final_output_type == "text":
            request_id = output.request_id
            text_output = output.outputs[0].text
            out_txt = os.path.join(output_dir, f"{request_id}.txt")
            lines = [
                "Prompt:\n",
                str(output.prompt) + "\n",
                "vllm_text_output:\n",
                str(text_output).strip() + "\n",
            ]
            try:
                with open(out_txt, "w", encoding="utf-8") as f:
                    f.writelines(lines)
            except Exception as e:
                print(f"[Warn] Failed writing text file {out_txt}: {e}")
            print(f"Request ID: {request_id}, Text saved to {out_txt}")
        elif stage_outputs.final_output_type == "audio":
            request_id = output.request_id
            audio_tensor = output.outputs[0].multimodal_output["audio"]
            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")

            if isinstance(audio_tensor, list):
                import torch

                audio_tensor = torch.cat(
                    [(t if isinstance(t, torch.Tensor) else torch.tensor(t)).flatten() for t in audio_tensor]
                )
            audio_numpy = audio_tensor.float().detach().cpu().numpy()
            if audio_numpy.ndim > 1:
                audio_numpy = audio_numpy.flatten()

            # MiniCPM-o 4.5 Token2Wav emits 24 kHz mono.
            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
            print(f"Request ID: {request_id}, Saved audio to {output_wav}")

        processed_count += 1
        if profiler_enabled and processed_count >= total_requests:
            print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
            omni.stop_profile(stages=args.profiler_stages)
            print("[Info] Waiting 30s for workers to write trace files to disk...")
            time.sleep(30)
            print("[Info] Trace export wait time finished.")
    omni.close()


def parse_args():
    parser = TrackingArgumentParser(
        description="Offline inference demo for MiniCPM-o 4.5 (text / image / audio / video → text + speech)"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="openbmb/MiniCPM-o-4_5",
        help="Model name or path.",
    )
    parser.add_argument(
        "--query-type",
        "-q",
        type=str,
        default="text",
        choices=query_map.keys(),
        help="Query type.",
    )
    parser.add_argument(
        "--log-stats",
        action="store_true",
        default=False,
        help="Enable writing detailed statistics (default: disabled)",
    )
    parser.add_argument(
        "--stage-init-timeout",
        type=int,
        default=300,
        help="Timeout for initializing a single stage in seconds (default: 300)",
    )
    parser.add_argument(
        "--batch-timeout",
        type=int,
        default=5,
        help="Timeout for batching in seconds (default: 5)",
    )
    parser.add_argument(
        "--init-timeout",
        type=int,
        default=300,
        help="Timeout for initializing stages in seconds (default: 300)",
    )
    parser.add_argument(
        "--shm-threshold-bytes",
        type=int,
        default=65536,
        help="Threshold for using shared memory in bytes (default: 65536)",
    )
    parser.add_argument(
        "--output-wav",
        default="output_audio",
        help="[Deprecated] Output wav directory (use --output-dir).",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Output directory for text/audio results.",
    )
    parser.add_argument(
        "--num-prompts",
        type=int,
        default=1,
        help="Number of prompts to generate.",
    )
    parser.add_argument(
        "--txt-prompts",
        type=str,
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help="Path to a deploy YAML (default: auto-load minicpmo_4_5.yaml).",
    )
    parser.add_argument(
        "--video-path",
        "-v",
        type=str,
        default=None,
        help="Path to local video file. If not provided, uses default video asset.",
    )
    parser.add_argument(
        "--image-path",
        "-i",
        type=str,
        default=None,
        help="Path to local image file. If not provided, uses default image asset.",
    )
    parser.add_argument(
        "--audio-path",
        "-a",
        type=str,
        default=None,
        help="Path to local audio file. If not provided, uses default audio asset.",
    )
    parser.add_argument(
        "--num-frames",
        type=int,
        default=16,
        help="Number of frames to extract from video (default: 16).",
    )
    parser.add_argument(
        "--sampling-rate",
        type=int,
        default=16000,
        help="Sampling rate for audio loading (default: 16000).",
    )
    parser.add_argument(
        "--log-dir",
        type=str,
        default="logs",
        help="Log directory (default: logs).",
    )
    parser.add_argument(
        "--modalities",
        type=str,
        default=None,
        help="Output modalities (comma-separated), e.g. text or text,audio.",
    )
    parser.add_argument(
        "--py-generator",
        action="store_true",
        default=False,
        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
    )
    parser.add_argument(
        "--enable-profiler",
        action="store_true",
        default=False,
        help="Enables profiling when set.",
    )
    parser.add_argument(
        "--profiler-stages",
        type=int,
        nargs="*",
        default=None,
        help="List of stage IDs to profile. If not set, profiles all stages.",
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="auto",
        help="Model dtype (auto, half, float16, bfloat16, float, float32).",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    main(args)