"""Mage-VL Proactive Streaming Commentary Demo (Colab-compatible).

Plays a video segment-by-segment and uses the built-in StreamMind gate
(System 1, mamba-ssm based) to decide *when* to generate commentary,
then invokes the full Qwen3-4B decoder (System 2) only on triggered segments.

Paper: arXiv:2607.24904
Official reference: mage_vl/inference_streaming.py
"""
from __future__ import annotations

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

import numpy as np
import torch

# ---------------------------------------------------------------------------
# mamba-ssm check (required for StreamMind gate)
# ---------------------------------------------------------------------------
try:
    import mamba_ssm  # noqa: F401
    _HAS_MAMBA = True
except ImportError:
    _HAS_MAMBA = False


USER_PROMPT = "Please describe the video content in detail based on the provided information."


# ---------------------------------------------------------------------------
# Video segment extraction (ffmpeg subclip)
# ---------------------------------------------------------------------------

def extract_subclip(src: str, start: float, dur: float, out: Path) -> Path:
    if out.exists() and out.stat().st_size > 0:
        return out
    out.parent.mkdir(parents=True, exist_ok=True)
    tmp = out.with_name(f".tmp_{out.name}")
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error",
         "-ss", f"{max(0.0, start):.3f}", "-t", f"{dur:.3f}",
         "-i", src, "-c", "copy",
         "-avoid_negative_ts", "make_zero", str(tmp)],
        check=True, timeout=120, capture_output=True,
    )
    tmp.replace(out)
    return out


def build_segments(video_path: str, segment_sec: float) -> list[dict]:
    """Split video into non-overlapping segments; returns [{path, start, end}, ...]."""
    try:
        import cv2
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps
        cap.release()
    except Exception:
        # Fallback: get duration via ffprobe
        try:
            result = subprocess.run(
                ["ffprobe", "-v", "error", "-show_entries",
                 "format=duration", "-of", "default=noprint_wrappers=1:nokey=1",
                 video_path],
                capture_output=True, text=True, timeout=30,
            )
            duration = float(result.stdout.strip())
        except Exception:
            duration = 60.0  # default assumption

    cache_dir = Path(video_path).resolve().parent / ".mage_clips"
    cache_dir.mkdir(parents=True, exist_ok=True)

    segments = []
    t = 0.0
    idx = 0
    while t < duration:
        end = min(duration, t + segment_sec)
        clip_path = cache_dir / f"seg_{idx:04d}_{t:.1f}-{end:.1f}.mp4"
        try:
            extract_subclip(video_path, t, end - t, clip_path)
            segments.append({"path": str(clip_path), "start": t, "end": end})
        except Exception as e:
            print(f"  [warn] skip {t:.1f}-{end:.1f}s: {e}")
        t = end
        idx += 1
    return segments


# ---------------------------------------------------------------------------
# Gate probabilities (System 1)
# ---------------------------------------------------------------------------

def compute_gate_probs(model, segments: list[dict], processor,
                       backend: str, num_frames: int, cur_fps: float,
                       device: str, dtype) -> list[float]:
    """Run StreamMind gate over all segments; returns p(speak) per segment."""
    gate_inputs = []
    for seg in segments:
        msgs = [{"role": "user", "content": [
            {"type": "video"}, {"type": "text", "text": USER_PROMPT}
        ]}]
        prompt = processor.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True,
        )
        kwargs = {
            "text": [prompt],
            "videos": [seg["path"]],
            "video_backend": backend,
            "return_tensors": "pt",
            "padding": False,
        }
        if backend == "codec":
            kwargs.update(
                codec_config={"patch": 16, "max_pixels": 150000},
                max_pixels=150000,
            )
        else:
            kwargs.update(num_frames=num_frames, target_fps=cur_fps)
        try:
            inp = processor(**kwargs)
            gate_inputs.append({
                "pixel_values": inp["pixel_values"].to(device=device, dtype=dtype),
                "image_grid_thw": inp["image_grid_thw"].to(device),
                "patch_positions": inp.get("patch_positions"),
            })
        except Exception as e:
            print(f"  [warn] gate input failed: {e}")
            gate_inputs.append(None)

    valid = [g for g in gate_inputs if g is not None]
    probs: list[float] = []
    if valid and hasattr(model, "streammind_gate_forward_segments"):
        with torch.inference_mode():
            logits = model.streammind_gate_forward_segments(valid)[0]
        lengths = [int(g["image_grid_thw"][:, 0].sum()) for g in valid]
        boundaries = torch.tensor(lengths, device=logits.device).cumsum(0) - 1
        probs_full = torch.softmax(logits[boundaries].float(), dim=-1)[:, 1].tolist()
        pi = 0
        for g in gate_inputs:
            if g is not None:
                probs.append(probs_full[pi])
                pi += 1
            else:
                probs.append(0.0)
    else:
        # Fallback: no gate available, all segments get p=1.0
        probs = [1.0] * len(gate_inputs)
    return probs


# ---------------------------------------------------------------------------
# Generate commentary for a single segment (System 2)
# ---------------------------------------------------------------------------

def generate_segment(model, processor, seg: dict, backend: str,
                     num_frames: int, cur_fps: float,
                     max_new_tokens: int, device: str, dtype) -> str:
    msgs = [{"role": "user", "content": [
        {"type": "video"}, {"type": "text", "text": USER_PROMPT}
    ]}]
    prompt = processor.apply_chat_template(
        msgs, tokenize=False, add_generation_prompt=True,
    )
    kwargs = {
        "text": [prompt],
        "videos": [seg["path"]],
        "video_backend": backend,
        "return_tensors": "pt",
        "padding": False,
    }
    if backend == "codec":
        kwargs.update(
            codec_config={"patch": 16, "max_pixels": 150000},
            max_pixels=150000,
        )
    else:
        kwargs.update(num_frames=num_frames, target_fps=cur_fps)
    inputs = processor(**kwargs)
    inputs = {
        k: (v.to(device=device, dtype=dtype) if k == "pixel_values" else v.to(device))
        for k, v in inputs.items()
    }
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    new = out[0, inputs["input_ids"].shape[1]:]
    return processor.tokenizer.decode(new, skip_special_tokens=True).strip()


# ---------------------------------------------------------------------------
# Main demo
# ---------------------------------------------------------------------------

def run_demo(args):
    from transformers import AutoModelForCausalLM, AutoProcessor

    print(f"[1/4] Loading model: {args.checkpoint}")
    t0 = time.time()
    processor = AutoProcessor.from_pretrained(args.checkpoint, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        args.checkpoint,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16 if args.gpu else torch.float32,
        attn_implementation=None if not args.gpu else "flash_attention_2",
    )
    device = "cuda" if args.gpu and torch.cuda.is_available() else "cpu"
    model = model.to(device).eval()
    dtype = model.dtype
    print(f"      Loaded in {time.time()-t0:.1f}s on {device}")

    has_gate = hasattr(model, "streammind_gate_forward_segments")
    if not _HAS_MAMBA:
        print("  [!] mamba-ssm not installed — StreamMind gate unavailable.")
        print("      Install: pip install mamba-ssm   (requires CUDA)")
    if not has_gate:
        print("  [!] Checkpoint missing streammind_gate_forward_segments(). "
              "Running un-gated (all segments will generate).")

    # ------------------------------------------------------------------
    # Build segments
    # ------------------------------------------------------------------
    print(f"[2/4] Splitting video into {args.segment_sec}s segments ...")
    segments = build_segments(args.video, args.segment_sec)
    if not segments:
        print("No segments extracted."); return
    print(f"      {len(segments)} segments ({segments[-1]['end']:.1f}s total)")

    # ------------------------------------------------------------------
    # Gate probabilities (System 1)
    # ------------------------------------------------------------------
    print("[3/4] Running StreamMind gate (System 1) ...")
    gate_probs = compute_gate_probs(
        model, segments, processor, args.backend,
        args.num_frames, args.cur_fps, device, dtype,
    )
    for i, (seg, p) in enumerate(zip(segments, gate_probs)):
        action = "SPEAK" if p >= args.threshold else "silence"
        print(f"      [{seg['start']:.1f}-{seg['end']:.1f}s] p={p:.2f} -> {action}")

    # ------------------------------------------------------------------
    # Generate commentary for triggered segments (System 2)
    # ------------------------------------------------------------------
    print("[4/4] Generating commentary (System 2) ...")
    print("=" * 70)

    for i, (seg, p) in enumerate(zip(segments, gate_probs)):
        elapsed = time.time()
        if p < args.threshold:
            print(f"\n[{seg['start']:.1f}-{seg['end']:.1f}s] "
                  f"gate=SILENCE (p={p:.2f})")
            sys.stdout.flush()
            continue

        print(f"\n[{seg['start']:.1f}-{seg['end']:.1f}s] "
              f"gate=RESPONSE (p={p:.2f}) -> generating ...", end="", flush=True)
        try:
            text = generate_segment(
                model, processor, seg, args.backend,
                args.num_frames, args.cur_fps,
                args.max_tokens, device, dtype,
            )
            t_gen = time.time() - elapsed
            print(f" ({t_gen:.1f}s)")
            # Wrap long text for terminal readability
            for j in range(0, len(text), 100):
                print(f"    {text[j:j+100]}")
        except Exception as e:
            print(f" (FAILED: {e})")
        sys.stdout.flush()

    print("\n" + "=" * 70)
    print("[done]")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(
        description="Mage-VL proactive streaming commentary demo "
                    "(Colab-compatible, text-only output)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    ap.add_argument("--video", required=True, help="Path to input video (.mp4)")
    ap.add_argument("--checkpoint", default="microsoft/Mage-VL",
                    help="HuggingFace model ID or local path")
    ap.add_argument("--backend", choices=["codec", "frames"], default="codec",
                    help="'codec'=Mage-ViT codec-native; 'frames'=uniform sampling")
    ap.add_argument("--segment-sec", type=float, default=8.0,
                    help="Segment length in seconds")
    ap.add_argument("--num-frames", type=int, default=16,
                    help="Max frames per segment (frames backend)")
    ap.add_argument("--cur-fps", type=float, default=2.0,
                    help="Target FPS for frame sampling")
    ap.add_argument("--threshold", type=float, default=0.5,
                    help="StreamMind gate threshold (p >= threshold => speak)")
    ap.add_argument("--max-tokens", type=int, default=120,
                    help="Max new tokens per commentary generation")
    ap.add_argument("--gpu", action="store_true",
                    help="Use CUDA (recommended; enables bf16 + flash-attn)")
    args = ap.parse_args()

    if not Path(args.video).is_file():
        ap.error(f"Video not found: {args.video}")

    if args.gpu and not torch.cuda.is_available():
        print("[!] --gpu set but no CUDA. Running on CPU.")
        args.gpu = False

    run_demo(args)


if __name__ == "__main__":
    main()