"""
MiniCPM-o 4.5 主动式流媒体视频分析 — Colab 可执行脚本

功能：
  - 上传或提供 mp4 URL
  - 按 segment_sec 切分视频
  - 对每段调用 MiniCPM-o 4.5 流式推理
  - 系统提示词驱动模型在检测到关键节点时主动输出推荐/评论
  - 可选语音输出（TTS）

用法：
  python minicpmo_proactive_streaming.py \
    --video /content/your_video.mp4 \
    --prompt "观察视频中投放的衣物，在每次出现新衣物时给出洗涤/护理推荐"
"""
from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch

# ---------------------------------------------------------------------------
# 1. 安装依赖（Colab 环境自动跳过已安装的包）
# ---------------------------------------------------------------------------
print("[1/5] 检查并安装依赖 ...")
subprocess_required = ["ffmpeg"]
import subprocess
for cmd in subprocess_required:
    if not shutil.which(cmd):
        print(f"  正在安装 {cmd} ...")
        subprocess.run(["apt-get", "update", "-qq"], check=False, capture_output=True)
        subprocess.run(["apt-get", "install", "-y", "-qq", cmd], check=False, capture_output=True)

import shutil

DEPS = [
    "transformers==4.51.0",
    "accelerate",
    "torch>=2.3.0,<=2.8.0",
    "torchaudio<=2.8.0",
    "minicpmo-utils[all]>=1.0.5",
    "librosa",
    "soundfile",
    "cv2",
]
for pkg in DEPS:
    try:
        __import__(pkg.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].replace("-", "_"))
    except ImportError:
        print(f"  安装 {pkg} ...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", pkg])
print("  ✓ 依赖就绪")

# ---------------------------------------------------------------------------
# 2. 加载模型
# ---------------------------------------------------------------------------
print("[2/5] 加载 MiniCPM-o 4.5 ...")
from transformers import AutoModel

MODEL_ID = "openbmb/MiniCPM-o-4_5"
model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    attn_implementation="sdpa",
    torch_dtype=torch.bfloat16,
    init_vision=True,
    init_audio=True,
    init_tts=True,
)
model.eval().cuda()
model.init_tts()
print(f"  ✓ 模型已加载 ({sum(p.numel() for p in model.parameters())/1e9:.1f}B 参数)")

# ---------------------------------------------------------------------------
# 3. 下载/准备视频
# ---------------------------------------------------------------------------
print("[3/5] 准备视频 ...")
import cv2


def download_or_use(path_or_url: str, dst_dir: Path) -> str:
    """本地路径直接返回；URL 则下载到 dst_dir。"""
    p = Path(path_or_url)
    if p.exists():
        return str(p)
    out = dst_dir / p.name
    if out.exists():
        return str(out)
    print(f"  正在下载: {path_or_url}")
    subprocess.run(
        ["curl", "-L", "-o", str(out), path_or_url],
        check=True, capture_output=True,
    )
    return str(out)


def get_video_duration(video_path: str) -> float:
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    cap.release()
    return frames / fps


def build_segments(video_path: str, segment_sec: float) -> list[dict]:
    """将视频切分为等长片段，返回 [{start, end, path}, ...]"""
    duration = get_video_duration(video_path)
    cache = Path(video_path).resolve().parent / ".segments"
    cache.mkdir(exist_ok=True)
    segments = []
    t = 0.0
    idx = 0
    while t < duration - 0.1:
        end = min(duration, t + segment_sec)
        clip = cache / f"seg_{idx:04d}_{t:.1f}-{end:.1f}.mp4"
        if not clip.exists():
            subprocess.run(
                ["ffmpeg", "-y", "-loglevel", "error",
                 "-ss", f"{t:.3f}", "-t", f"{end - t:.3f}",
                 "-i", video_path, "-c", "copy",
                 "-avoid_negative_ts", "make_zero", str(clip)],
                check=True, capture_output=True,
            )
        segments.append({"path": str(clip), "start": t, "end": end})
        t = end
        idx += 1
    return segments


# ---------------------------------------------------------------------------
# 4. 流式推理主循环
# ---------------------------------------------------------------------------
print("[4/5] 开始流式推理 ...")
from minicpmo.utils import get_video_frame_audio_segments

DEFAULT_PROMPT = (
    "你是一位专业的衣物护理顾问。请观看视频，"
    "在每次检测到新的衣物类型被投放/展示时，"
    "主动给出简短的洗涤/护理推荐（中英文均可）。"
    "如果画面中没有新衣物出现，保持沉默不输出。"
    "输出格式：[时间戳] 推荐内容"
)


def run_segment(model, session_id: str, seg: dict,
                sys_msg: dict, user_prompt: str,
                generate_audio: bool = False) -> dict:
    """对单个视频片段运行 streaming_prefill + streaming_generate，返回 {text, audio_path, start, end}"""
    seg_path = seg["path"]
    start = seg["start"]

    # 提取帧和音频段
    try:
        video_frames, audio_segments, stacked_frames = get_video_frame_audio_segments(
            seg_path, stack_frames=1, use_ffmpeg=True, adjust_audio_length=True
        )
    except Exception as e:
        print(f"  [!] 视频处理失败 {seg_path}: {e}")
        return {"start": start, "end": seg["end"], "text": f"[错误] {e}", "audio_path": None}

    if not video_frames:
        return {"start": start, "end": seg["end"], "text": "[无帧]", "audio_path": None}

    # Step 1: Prefill 系统提示
    model.streaming_prefill(session_id=session_id, msgs=[sys_msg])

    # Step 2: Prefill 视频帧 + 音频
    omni_contents = []
    for i, frame in enumerate(video_frames):
        omni_contents.append(frame)
        if i < len(audio_segments):
            omni_contents.append(audio_segments[i])
        if stacked_frames is not None and i < len(stacked_frames) and stacked_frames[i] is not None:
            omni_contents.append(stacked_frames[i])
    # 追加用户指令
    omni_contents.append(user_prompt)

    last_idx = len(omni_contents) - 1
    for idx, content in enumerate(omni_contents):
        is_last = (idx == last_idx)
        model.streaming_prefill(
            session_id=session_id,
            msgs=[{"role": "user", "content": [content]}],
            omni_mode=True,
            is_last_chunk=is_last,
        )

    # Step 3: 生成回复
    audio_chunks = []
    text_parts = []
    try:
        gen = model.streaming_generate(
            session_id=session_id,
            generate_audio=generate_audio,
            use_tts_template=True,
            enable_thinking=False,
            do_sample=True,
            temperature=0.7,
            max_new_tokens=256,
        )
        for item in gen:
            if generate_audio and isinstance(item, tuple) and len(item) == 2:
                wav_chunk, text_chunk = item
                audio_chunks.append(wav_chunk)
                text_parts.append(text_chunk)
            else:
                text_chunk, is_finished = item
                text_parts.append(text_chunk)
    except Exception as e:
        return {"start": start, "end": seg["end"], "text": f"[生成错误] {e}", "audio_path": None}

    text = "".join(text_parts).strip()
    audio_path = None
    if generate_audio and audio_chunks:
        generated_waveform = torch.cat(audio_chunks, dim=-1)[0]
        audio_path = f"/content/proactive_output_{int(start):04d}.wav"
        import soundfile as sf
        sf.write(audio_path, generated_waveform.cpu().numpy(), samplerate=24000)

    return {"start": start, "end": seg["end"], "text": text, "audio_path": audio_path}


# ---------------------------------------------------------------------------
# 5. 主入口
# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description="MiniCPM-o 4.5 主动式流媒体视频分析")
    ap.add_argument("--video", required=True, help="mp4 本地路径 或 HTTP URL")
    ap.add_argument("--prompt", default=DEFAULT_PROMPT, help="任务指令 prompt")
    ap.add_argument("--language", default="zh", choices=["zh", "en"], help="系统提示语言")
    ap.add_argument("--segment-sec", type=float, default=10.0, help="每段时长（秒）")
    ap.add_argument("--threshold", type=float, default=0.0,
                    help="仅当检测到新内容时才输出（当前版本始终输出）")
    ap.add_argument("--output-audio", action="store_true", help="是否生成语音输出")
    ap.add_argument("--ref-audio", default=None, help="参考音频路径（用于声音克隆，可选）")
    args = ap.parse_args()

    # 下载/定位视频
    video_path = download_or_use(args.video, Path("/content"))
    print(f"  视频: {video_path} ({get_video_duration(video_path):.1f}s)")

    # 构建片段列表
    segments = build_segments(video_path, args.segment_sec)
    print(f"  共 {len(segments)} 个片段 (每段 {args.segment_sec}s)")

    # 系统提示
    if args.language == "zh":
        sys_content = [
            "你是一个智能衣物护理助手，持续观看视频流。",
            "当画面中出现新的衣物/织物时，主动给出简短的洗涤或护理推荐。",
            "若未检测到新内容则不要输出。",
            "每次输出以 [MM:SS] 开头标注时间戳。",
        ]
    else:
        sys_content = [
            "You are a smart clothing care assistant watching a video stream.",
            "When new clothing/fabric appears, proactively give a brief washing or care recommendation.",
            "Stay silent if nothing new is detected.",
            "Start each output with a [MM:SS] timestamp.",
        ]
    sys_msg = {"role": "system", "content": sys_content}

    # 可选：声音克隆
    if args.ref_audio and Path(args.ref_audio).exists():
        import librosa
        ref_audio, _ = librosa.load(args.ref_audio, sr=16000, mono=True)
        model.init_token2wav_cache(ref_audio)
        print(f"  声音克隆参考音频已加载: {args.ref_audio}")

    session_id = "proactive_session"
    model.reset_session()

    # 运行推理
    print("[5/5] 推理中 ...")
    print("=" * 70)
    all_results = []
    t_total = time.time()

    for seg in segments:
        r = run_segment(
            model, session_id, seg,
            sys_msg=sys_msg,
            user_prompt=args.prompt,
            generate_audio=args.output_audio,
        )
        all_results.append(r)
        elapsed = time.time() - t_total
        status = "🔊" if r["audio_path"] else "📝"
        print(f"[{status}] [{r['start']:.1f}-{r['end']:.1f}s] "
              f"{r['text'][:120] or '(沉默)'}  ({elapsed:.1f}s)")
        sys.stdout.flush()

    print("=" * 70)
    print(f"✓ 完成！总耗时 {time.time()-t_total:.1f}s，共 {len(all_results)} 段")

    # 汇总输出
    print("\n=== 关键时间节点推荐汇总 ===")
    for r in all_results:
        if r["text"] and r["text"].strip() and r["text"].strip() not in ("(沉默)", "[无帧]"):
            print(f"  [{r['start']:.1f}s] {r['text']}")

    # 合并音频（如果有）
    if args.output_audio:
        audio_files = [r["audio_path"] for r in all_results if r["audio_path"]]
        if audio_files:
            import soundfile as sf
            merged = np.concatenate([sf.read(f, dtype="float32")[0] for f in audio_files])
            out_audio = "/content/proactive_full_output.wav"
            sf.write(out_audio, merged, samplerate=24000)
            print(f"\n完整语音输出: {out_audio}")

    return all_results


if __name__ == "__main__":
    main()