# ============================================================
# Mage-VL 第五版
# T4 16GB / 足球视频 / 4秒一段 / 真正逐段实时流式分析
#
# 核心流程：
#
# 视频
#   ↓
# 当前 4 秒
#   ↓
# 精确切片（重新编码）
#   ↓
# Codec
#   ↓
# Gate
#   ↓
# ┌───────────────┐
# │ p < threshold │ → 跳过
# │ p >= threshold│ → Mage-VL
# └───────────────┘
#   ↓
# 输出
#   ↓
# 清理 GPU
#   ↓
# 下一 4 秒
#
# 不会提前 Codec 全部视频
# ============================================================

import os

# ------------------------------------------------------------
# 0. 环境变量
# ------------------------------------------------------------

os.environ.setdefault(
    "PYTORCH_CUDA_ALLOC_CONF",
    "expandable_segments:True,max_split_size_mb:128"
)
os.environ.setdefault(
    "ONLINE_CODEC_CACHE_DIR",
    "/tmp/mage_vl_codec_cache"
)
os.environ.setdefault(
    "STREAMMIND_CLIP_CACHE",
    "/tmp/mage_vl_clips"
)
os.makedirs(
    os.environ["ONLINE_CODEC_CACHE_DIR"],
    exist_ok=True
)
os.makedirs(
    os.environ["STREAMMIND_CLIP_CACHE"],
    exist_ok=True
)

# ------------------------------------------------------------
# 1. Python imports
# ------------------------------------------------------------

import gc
import time
import warnings
import subprocess
import importlib
from pathlib import Path

import torch

from transformers import (
    AutoModelForCausalLM,
    AutoProcessor,
)

# ------------------------------------------------------------
# 2. 基本配置
# ------------------------------------------------------------

MODEL_ID = "microsoft/Mage-VL"
VIDEO_FILE = "/content/soccer-broadcast.mp4"
# 每次分析多少秒
SEGMENT_SEC = 4.0
# 最多分析多少段
#
# None = 整个视频
#
# 如果只是测试，可以设置成 8
MAX_SEGMENTS = None
# Gate 阈值
GATE_THRESHOLD = 0.50
# 最大生成 token
MAX_NEW_TOKENS = 64
# Codec 最大像素
#
# T4 16GB 推荐先保持 64000
#
# 如果显存仍然紧张：
# 48000
# 32000
#
CODEC_MAX_PIXELS = 64_000
# Codec target canvas
CODEC_TARGET_CANVAS = 24
# 是否模拟真实时间
#
# True：
# 如果处理速度 > 1x，则等待
#
# False：
# 处理完马上进入下一段
SIMULATE_REALTIME = True
# 为了避免频繁 empty_cache
# 可以每 N 个 segment 强制一次
EMPTY_CACHE_EVERY = 1

# ------------------------------------------------------------
# 3. Prompt
# ------------------------------------------------------------

STREAM_PROMPT = """
请用中文分析这段足球视频。
重点描述：
1. 视频中有哪些人物或球员
2. 球员正在做什么
3. 当前比赛发生了什么
4. 是否有射门、传球、带球、抢断、犯规、角球、任意球、进球等动作
5. 如果可以判断，请说明球员的大致位置和动作
6. 判断当前是什么体育项目
只描述视频中能够观察到的内容，不要猜测无法确认的信息。
尽量简洁，适合实时视频分析。
""".strip()

# ------------------------------------------------------------
# 4. CUDA / dtype
# ------------------------------------------------------------

if not torch.cuda.is_available():
    raise RuntimeError(
        "没有检测到 CUDA。请确认当前环境已经启用 NVIDIA GPU。"
    )
DEVICE = "cuda"
GPU_NAME = torch.cuda.get_device_name(0)
print("=" * 70)
print("Mage-VL 第五版")
print("=" * 70)
print(f"[GPU] {GPU_NAME}")
print(f"[CUDA] {torch.version.cuda}")
print(f"[PyTorch] {torch.__version__}")
print("=" * 70)

# ------------------------------------------------------------
# 5. T4 使用 FP16
# ------------------------------------------------------------
#
# T4 是 Turing 架构。
#
# 相比 BF16：
# T4 上 FP16 通常更加合适。
#
# ------------------------------------------------------------

MODEL_DTYPE = torch.float16

# ------------------------------------------------------------
# 6. 加载 Processor
# ------------------------------------------------------------

print()
print("[init] 正在加载 Processor...")
processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)
print("[init] Processor 加载完成")

# ------------------------------------------------------------
# 7. 加载模型
# ------------------------------------------------------------

print("[init] 正在加载 Mage-VL...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=MODEL_DTYPE,
    attn_implementation="sdpa",
)
model = model.to(DEVICE).eval()
print("[init] Mage-VL 加载完成")

# ------------------------------------------------------------
# 8. 找到远程 codec 模块
# ------------------------------------------------------------

_REMOTE_PKG = type(processor).__module__.rsplit(".", 1)[0]
codec_mod = importlib.import_module(
    _REMOTE_PKG + ".codec_video_processing_mage_vl"
)
CodecConfig = codec_mod.CodecConfig

# ------------------------------------------------------------
# 9. 加载 Cognition Gate
# ------------------------------------------------------------

try:

    model.model._load_streammind_gate()
    print("[gate] 认知门控加载成功")
except Exception as exc:

    print(
        f"[gate] 认知门控加载失败: {type(exc).__name__}: {exc}"
    )

# ------------------------------------------------------------
# 10. 显存信息
# ------------------------------------------------------------

def print_gpu_memory(prefix="GPU"):

    if not torch.cuda.is_available():
        return
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    total = torch.cuda.get_device_properties(0).total_memory
    total = total / 1024**3
    free = total - allocated
    print(
        f"[{prefix}] "
        f"allocated={allocated:.2f}GB "
        f"reserved={reserved:.2f}GB "
        f"free≈{free:.2f}GB "
        f"total={total:.2f}GB",
        flush=True,
    )

# ------------------------------------------------------------
# 11. FFprobe 获取视频时长
# ------------------------------------------------------------

def _ffprobe_duration(video_path):

    try:

        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=nw=1:nk=1",
                str(video_path),
            ],
            check=True,
            capture_output=True,
            timeout=60,
            text=True,
        )
        value = result.stdout.strip()
        if not value:
            return 0.0
        return float(value)
    except Exception as exc:

        print(
            f"[ffprobe] 获取视频时长失败: {exc}",
            flush=True,
        )
        return 0.0

# ------------------------------------------------------------
# 12. Codec 配置
# ------------------------------------------------------------

def _codec_cfg(
    target_canvas=CODEC_TARGET_CANVAS,
    max_pixels=CODEC_MAX_PIXELS,
):
    override = {
        "engine": "hevc",
        "target_canvas": int(target_canvas),
        "patch": 16,
    }
    kwargs = dict(
        processor._codec_config_defaults
    )
    merged_dcvc = dict(
        kwargs.get("dcvc") or {}
    )
    merged_dcvc.update(
        override.get("dcvc") or {}
    )
    kwargs.update(override)
    if merged_dcvc:
        kwargs["dcvc"] = merged_dcvc
    kwargs["max_pixels"] = int(max_pixels)
    return CodecConfig(**kwargs), override

# ------------------------------------------------------------
# 13. Prompt
# ------------------------------------------------------------

def _prompt(kind, question):

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": kind
                },
                {
                    "type": "text",
                    "text": question,
                },
            ],
        }
    ]
    return processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

# ------------------------------------------------------------
# 14. CPU → GPU
# ------------------------------------------------------------

def _to_cuda(inputs):

    output = {}
    for key, value in inputs.items():
        if not hasattr(value, "to"):
            continue
        value = value.to(DEVICE)
        if key == "pixel_values":
            value = value.to(MODEL_DTYPE)
        output[key] = value
    return output

# ------------------------------------------------------------
# 15. 精确切片
# ------------------------------------------------------------
#
# 重要：
#
# 原来的：
#
#   -c copy
#
# 是直接复制码流。
#
# 这样 -ss 可能落在关键帧之外，
# 导致实际得到的视频开始时间不够精确。
#
# 这里改成重新编码。
#
# 优点：
#   时间窗口更加准确
#
# 缺点：
#   CPU 会多一点开销。
#
# 对“模拟实时分析”而言，
# 这个方案更稳。
#
# ------------------------------------------------------------

def _extract_subclip(
    src,
    start,
    duration,
    out_path,
):
    if duration <= 0:
        raise ValueError(
            f"duration 必须 > 0，当前={duration}"
        )
    tmp = out_path.with_name(
        f".tmp_{out_path.name}"
    )
    if tmp.exists():
        try:
            tmp.unlink()
        except Exception:
            pass
    command = [
        "ffmpeg",
        "-y",
        "-loglevel",
        "error",
        # 精确时间定位
        "-ss",
        f"{max(0.0, start):.3f}",
        "-i",
        str(src),
        "-t",
        f"{duration:.3f}",
        # 视频重新编码
        "-c:v",
        "libx264",
        # 很快的 CPU 编码
        "-preset",
        "ultrafast",
        # 保证足够好的视觉质量
        "-crf",
        "23",
        # 避免音频处理
        "-an",
        # 独立可解码
        "-pix_fmt",
        "yuv420p",
        str(tmp),
    ]
    subprocess.run(
        command,
        check=True,
        capture_output=True,
        timeout=180,
    )
    if not tmp.exists() or tmp.stat().st_size <= 0:
        raise RuntimeError(
            f"FFmpeg 没有生成有效文件: {tmp}"
        )
    if out_path.exists():
        try:
            out_path.unlink()
        except Exception:
            pass
    tmp.replace(out_path)
    return out_path

# ------------------------------------------------------------
# 16. 清理显存
# ------------------------------------------------------------

def _cleanup_cuda(*objects):

    for obj in objects:
        try:
            del obj
        except Exception:
            pass
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass

# ------------------------------------------------------------
# 17. Gate
# ------------------------------------------------------------

@torch.inference_mode()
def _run_gate(cpu_inputs):

    """
    只把当前 segment 的 Gate 所需视觉特征送入 GPU。
    不保存多个 segment。
    """

    cuda_visual_features = {}
    for key in (
        "pixel_values",
        "image_grid_thw",
        "patch_positions",
    ):
        if key not in cpu_inputs:
            continue
        value = cpu_inputs[key].to(DEVICE)
        if key == "pixel_values":
            value = value.to(MODEL_DTYPE)
        cuda_visual_features[key] = value
    try:

        logits = (
            model.model.streammind_gate_forward_segments(
                [cuda_visual_features]
            )[0]
        )
        # 当前只有一个 segment
        #
        # image_grid_thw 第一维是 segment/frame 相关结构，
        # 与你第四版保持同样的计算方式。

        length = int(
            cpu_inputs["image_grid_thw"][:, 0].sum()
        )
        boundary = length - 1
        probability = torch.softmax(
            logits[boundary].float(),
            dim=-1,
        )[1].item()
        return float(probability)
    finally:

        _cleanup_cuda(
            cuda_visual_features,
            locals().get("logits"),
        )

# ------------------------------------------------------------
# 18. Mage-VL 生成
# ------------------------------------------------------------

@torch.inference_mode()
def _generate(
    inputs,
    max_new_tokens=MAX_NEW_TOKENS,
):
    output = model.generate(
        **inputs,
        max_new_tokens=int(
            max_new_tokens
        ),
        do_sample=False,
        use_cache=True,
    )
    prompt_length = inputs[
        "input_ids"
    ].shape[1]
    generated = output[
        0,
        prompt_length:
    ]
    text = processor.tokenizer.decode(
        generated,
        skip_special_tokens=True,
    ).strip()
    del output
    del generated
    return text

# ------------------------------------------------------------
# 19. 计算实时倍率
# ------------------------------------------------------------

def _realtime_speed(
    video_duration,
    processing_time,
):
    if processing_time <= 0:
        return 0.0
    return (
        video_duration /
        processing_time
    )

# ------------------------------------------------------------
# 20. 单 segment 分析
# ------------------------------------------------------------

def _analyze_segment(
    clip,
    text,
    codec_override,
):

    # --------------------------------------------------------
    # Codec
    # --------------------------------------------------------

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        inputs = processor(
            text=[text],
            videos=[
                str(clip)
            ],
            video_backend="codec",
            max_pixels=CODEC_MAX_PIXELS,
            codec_config=codec_override,
            return_tensors="pt",
            padding=False,
        )

    # --------------------------------------------------------
    # Gate
    # --------------------------------------------------------

    probability = _run_gate(
        inputs
    )
    return inputs, probability

# ------------------------------------------------------------
# 21. 真正流式分析
# ------------------------------------------------------------

def stream_video_analysis(
    video_path: str,
    segment_sec: float = SEGMENT_SEC,
    gate_threshold: float = GATE_THRESHOLD,
    max_segments=MAX_SEGMENTS,
    max_new_tokens: int = MAX_NEW_TOKENS,
    simulate_realtime: bool = SIMULATE_REALTIME,
):
    path = Path(video_path)
    if not path.exists():
        raise FileNotFoundError(
            f"找不到视频文件: {video_path}"
        )

    # --------------------------------------------------------
    # 获取时长
    # --------------------------------------------------------

    duration = _ffprobe_duration(path)
    if duration <= 0:
        raise RuntimeError(
            "无法获得视频时长。"
        )

    # --------------------------------------------------------
    # 临时目录
    # --------------------------------------------------------

    clip_dir = Path(
        os.environ["STREAMMIND_CLIP_CACHE"]
    )
    clip_dir.mkdir(
        parents=True,
        exist_ok=True,
    )

    # --------------------------------------------------------
    # Prompt
    # --------------------------------------------------------

    text = _prompt(
        "video",
        STREAM_PROMPT,
    )

    # --------------------------------------------------------
    # Codec 配置
    # --------------------------------------------------------

    _, codec_override = _codec_cfg(
        CODEC_TARGET_CANVAS
    )

    # --------------------------------------------------------
    # 状态
    # --------------------------------------------------------

    start = 0.0
    index = 0
    total_processed = 0.0
    all_lines = []
    yield (
        f"开始实时流式分析\n"
        f"视频总时长: {duration:.1f}s\n"
        f"分析窗口: {segment_sec:.1f}s\n"
        f"Gate threshold: {gate_threshold:.2f}\n"
        f"模拟实时: {simulate_realtime}"
    )
    # ========================================================
    # 核心循环
    # ========================================================

    while start < duration - 0.05:
        if (
            max_segments is not None
            and index >= max_segments
        ):
            break

        # ----------------------------------------------------
        # 当前 segment
        # ----------------------------------------------------

        end = min(
            duration,
            start + segment_sec,
        )
        current_duration = end - start
        index += 1
        segment_begin_time = time.time()
        clip = clip_dir / (
            f"{path.stem}"
            f"__segment_{index:06d}"
            f"__{start:.3f}_{end:.3f}.mp4"
        )

        # ----------------------------------------------------
        # 状态输出
        # ----------------------------------------------------

        yield (
            "\n".join(all_lines)
            + (
                "\n\n"
                if all_lines
                else ""
            )
            + f"⏱️ [{index}] "
            f"正在分析 "
            f"{start:.1f}–{end:.1f}s ..."
        )
        inputs = None
        cuda_inputs = None
        try:

            # =================================================
            # STEP 1
            # 精确切片
            # =================================================

            extract_begin = time.time()
            _extract_subclip(
                path,
                start,
                current_duration,
                clip,
            )
            extract_time = (
                time.time()
                - extract_begin
            )
            # =================================================
            # STEP 2
            # Codec + Gate
            # =================================================

            codec_begin = time.time()
            inputs, probability = (
                _analyze_segment(
                    clip,
                    text,
                    codec_override,
                )
            )
            codec_gate_time = (
                time.time()
                - codec_begin
            )
            # =================================================
            # Gate 结果
            # =================================================

            head = (
                f"**t = {start:.1f}–{end:.1f}s** "
                f"· Gate = `{probability:.2f}`"
            )

            # -------------------------------------------------
            # Gate 不通过
            # -------------------------------------------------

            if probability < float(
                gate_threshold
            ):
                processing_time = (
                    time.time()
                    - segment_begin_time
                )
                speed = _realtime_speed(
                    current_duration,
                    processing_time,
                )
                line = (
                    f"{head} → 🔇 "
                    f"*跳过生成* "
                    f"(processing={processing_time:.2f}s, "
                    f"realtime={speed:.2f}×)"
                )
                all_lines.append(line)
                yield (
                    "\n\n".join(all_lines)
                )

            # -------------------------------------------------
            # Gate 通过
            # -------------------------------------------------

            else:
                yield (
                    "\n\n".join(all_lines)
                    + (
                        "\n\n"
                        if all_lines
                        else ""
                    )
                    + (
                        f"{head} → 🧠 "
                        f"*正在生成...*"
                    )
                )
                # =============================================
                # STEP 3
                # 当前 segment 才进入 LLM
                # =============================================

                generate_begin = time.time()
                cuda_inputs = _to_cuda(
                    inputs
                )
                said = _generate(
                    cuda_inputs,
                    max_new_tokens,
                )
                generate_time = (
                    time.time()
                    - generate_begin
                )
                processing_time = (
                    time.time()
                    - segment_begin_time
                )
                speed = _realtime_speed(
                    current_duration,
                    processing_time,
                )
                line = (
                    f"{head} → 🗣️ {said}\n"
                    f"   Codec/Gate="
                    f"{codec_gate_time:.2f}s"
                    f" · Generate="
                    f"{generate_time:.2f}s"
                    f" · Total="
                    f"{processing_time:.2f}s"
                    f" · realtime="
                    f"{speed:.2f}×"
                )
                all_lines.append(line)
                yield (
                    "\n\n".join(all_lines)
                )

            # -------------------------------------------------
            # 统计
            # -------------------------------------------------

            processing_time = (
                time.time()
                - segment_begin_time
            )
            total_processed += (
                processing_time
            )

            # -------------------------------------------------
            # 模拟实时
            #
            # 如果：
            #
            # 4 秒视频
            # 实际只花 2 秒
            #
            # 等 2 秒再进入下一段。
            #
            # 如果实际用了 6 秒：
            #
            # 不等待。
            #
            # 这样不会越来越落后。
            # -------------------------------------------------

            if simulate_realtime:
                remaining = (
                    current_duration
                    - processing_time
                )
                if remaining > 0:
                    yield (
                        "\n\n".join(all_lines)
                        + "\n\n"
                        + (
                            f"⏳ 当前处理速度 "
                            f"{speed:.2f}×，"
                            f"等待 {remaining:.2f}s "
                            f"模拟实时..."
                        )
                    )
                    time.sleep(
                        remaining
                    )
        except Exception as exc:

            processing_time = (
                time.time()
                - segment_begin_time
            )
            line = (
                f"**t = {start:.1f}–{end:.1f}s** "
                f"→ ❌ "
                f"{type(exc).__name__}: "
                f"{exc}"
            )
            all_lines.append(line)
            yield (
                "\n\n".join(all_lines)
            )
        finally:

            # =================================================
            # STEP 4
            # 清理当前 segment
            # =================================================

            _cleanup_cuda(
                cuda_inputs,
                inputs,
            )
            cuda_inputs = None
            inputs = None

            # ------------------------------------------------
            # 删除临时视频
            # ------------------------------------------------

            try:

                if clip.exists():
                    clip.unlink()
            except Exception:
                pass

            # ------------------------------------------------
            # 定期清理 GPU
            # ------------------------------------------------

            if (
                index % EMPTY_CACHE_EVERY
                == 0
            ):
                gc.collect()
                torch.cuda.empty_cache()

            # ------------------------------------------------
            # 下一段
            # ------------------------------------------------

            start = end

            # ------------------------------------------------
            # 显存监控
            # ------------------------------------------------

            print_gpu_memory(
                prefix=f"segment {index}"
            )
    # ========================================================
    # 完成
    # ========================================================

    yield (
        "\n\n".join(all_lines)
        + "\n\n"
        + "=" * 60
        + "\n"
        + "✅ 实时流式分析完成"
        + "\n"
        + f"共处理 {index} 个 segment"
        + "\n"
        + f"视频时长 {min(start, duration):.1f}s"
        + "\n"
        + f"累计处理时间 {total_processed:.1f}s"
        + "\n"
        + "=" * 60
    )
# ============================================================
# 22. Main
# ============================================================

if __name__ == "__main__":
    print()
    print("=" * 70)
    print("开始测试 Mage-VL 实时流式视频分析")
    print("=" * 70)
    print()
    for chunk_output in stream_video_analysis(
        video_path=VIDEO_FILE,
        segment_sec=4.0,
        gate_threshold=0.50,
        max_segments=MAX_SEGMENTS,
        max_new_tokens=64,
        simulate_realtime=True,
    ):
        # 清屏
        print(
            "\033[H\033[J",
            end=""
        )
        print(
            "=== Mage-VL 实时流式分析 ==="
        )
        print()
        print(chunk_output)
        print()
        print("-" * 70)