###第四版，优化可在T416G跑起来，模拟实时分析
import os
import gc
import time
import warnings
import subprocess
import importlib
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoProcessor

# ==========================================
# 1. 环境变量与 CUDA 分配策略优化
# ==========================================
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)

MODEL_ID = "microsoft/Mage-VL"
CODEC_MAX_PIXELS = 64_000  # 优化1：稍微降低分辨率上限，显著减少 Token 显存占用
STREAM_PROMPT = "Please describe the video content in detail based on the provided information."

print("[init] 加载模型中...")
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()

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

# 加载 Cognition Gate
try:
    model.model._load_streammind_gate()
    print("[gate] 认知门控加载成功")
except Exception as exc:
    print(f"[gate] 加载失败: {exc}")

# ==========================================
# 2. 辅助工具函数
# ==========================================
def _ffprobe_duration(video_path):
    try:
        out = 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,
        )
        return float(out.stdout.strip())
    except Exception:
        return 0.0

def _codec_cfg(target_canvas=12, max_pixels=CODEC_MAX_PIXELS): # 优化2：target_canvas 从 32 降到 24
    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

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)

def _to_cuda(inputs):
    out = {}
    for key, value in inputs.items():
        if not hasattr(value, "to"):
            continue
        out[key] = value.to("cuda")
        if key == "pixel_values":
            out[key] = out[key].to(model.dtype)
    return out

def _extract_subclip(src, start, duration, out_path):
    if out_path.exists() and out_path.stat().st_size > 0:
        return out_path
    tmp = out_path.with_name(f".tmp_{out_path.name}")
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-ss", f"{max(0.0, start):.3f}",
         "-t", f"{duration:.3f}", "-i", str(src), "-c", "copy",
         "-avoid_negative_ts", "make_zero", str(tmp)],
        check=True, capture_output=True, timeout=180,
    )
    tmp.replace(out_path)
    return out_path

@torch.inference_mode()
def _generate(inputs, max_new_tokens):
    # 清理缓存以保证生成拥有最大连续显存块
    gc.collect()
    torch.cuda.empty_cache()
    
    out = model.generate(
        **inputs, 
        max_new_tokens=int(max_new_tokens), 
        do_sample=False,
        use_cache=True
    )
    new = out[0, inputs["input_ids"].shape[1]:]
    res = processor.tokenizer.decode(new, skip_special_tokens=True).strip()
    
    # 手动释放临时缓存
    del out, new
    gc.collect()
    torch.cuda.empty_cache()
    return res

# ==========================================
# 3. 优化后的流式分析函数
# ==========================================
def stream_video_analysis(
    video_path: str,
    segment_sec: float = 4.0,
    gate_threshold: float = 0.5,
    max_segments: int = 8,
    max_new_tokens: int = 64  # 优化3：适度缩减 token 生成长度，流式推理推荐 48-64
):
    path = Path(video_path)
    if not path.exists():
        raise FileNotFoundError(f"找不到视频文件: {video_path}")

    duration = _ffprobe_duration(path)
    clip_dir = Path("/tmp/mage_vl_clips")
    clip_dir.mkdir(parents=True, exist_ok=True)
    stem = path.stem
    text = _prompt("video", STREAM_PROMPT)
    cfg, override = _codec_cfg(24)

    raw_cpu_segments, spans = [], []
    start = 0.0
    index = 0

    yield f"开始切片视频 (总时长: {duration:.1f}s)..."

    # 1. 切片与预处理（保留在 CPU，避免占用 GPU 显存）
    while start < duration - 0.2 and index < max_segments:
        end = min(duration, start + segment_sec)
        clip = clip_dir / f"{stem}__t{int(start * 100):08d}__d{int((end - start) * 100):08d}.mp4"
        
        yield f"正在进行第 {index + 1} 个片段的 Codec 分析 (t={start:.1f}–{end:.1f}s)..."
        try:
            _extract_subclip(path, start, end - start, clip)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                inputs = processor(
                    text=[text], videos=[str(clip)], video_backend="codec",
                    max_pixels=CODEC_MAX_PIXELS, codec_config=override,
                    return_tensors="pt", padding=False,
                )
        except Exception as exc:
            print(f"[stream] segment {start:.1f}-{end:.1f} unusable: {exc}", flush=True)
            start, index = end, index + 1
            continue

        raw_cpu_segments.append(inputs)
        spans.append((start, end))
        start, index = end, index + 1

    if not raw_cpu_segments:
        yield "未检测到可用片段。"
        return

    # 2. 评估阶段：仅将 Gate 关注的轻量级特征输入 GPU 计算分数
    yield f"正在使用认知门控 (System 1) 评估 {len(raw_cpu_segments)} 个片段..."
    
    cuda_visual_features = []
    for seg in raw_cpu_segments:
        feat = {}
        for k in ("pixel_values", "image_grid_thw", "patch_positions"):
            if k in seg:
                feat[k] = seg[k].to("cuda")
                if k == "pixel_values":
                    feat[k] = feat[k].to(model.dtype)
        cuda_visual_features.append(feat)

    with torch.inference_mode():
        logits = model.model.streammind_gate_forward_segments(cuda_visual_features)[0]
    
    lengths = [int(seg["image_grid_thw"][:, 0].sum()) for seg in cuda_visual_features]
    boundaries = torch.tensor(lengths, device=logits.device).cumsum(0) - 1
    probs = torch.softmax(logits[boundaries].float(), dim=-1)[:, 1].tolist()

    # 释放 gate 特征占用的显存
    del cuda_visual_features, logits, boundaries
    gc.collect()
    torch.cuda.empty_cache()

    # 3. 按片段迭代生成并按需按时传输至 GPU
    lines = []
    for (span, cpu_inputs, prob) in zip(spans, raw_cpu_segments, probs):
        head = f"**t = {span[0]:.1f}–{span[1]:.1f}s** · p(respond) = `{prob:.2f}`"
        
        if prob < float(gate_threshold):
            lines.append(f"{head} → 🔇 *silence*")
            yield "\n\n".join(lines)
            continue
        
        lines.append(f"{head} → 🗣️ *generating…*")
        yield "\n\n".join(lines)
        
        try:
            # 仅在真正需要生成时，将当前片段输入移至 GPU
            cuda_inputs = _to_cuda(cpu_inputs)
            said = _generate(cuda_inputs, max_new_tokens)
            del cuda_inputs
        except Exception as exc:
            said = f"(generation failed: {type(exc).__name__}: {exc})"
            
        lines[-1] = f"{head} → 🗣️ {said}"
        yield "\n\n".join(lines)

    yield "\n\n".join(lines) + f"\n\n> 处理完成。"

# ==========================================
# 4. 执行测试
# ==========================================
if __name__ == "__main__":
    VIDEO_FILE = "/content/soccer-broadcast.mp4"
    
    for chunk_output in stream_video_analysis(
        video_path=VIDEO_FILE,
        segment_sec=4.0,
        gate_threshold=0.5,
        max_segments=8,
        max_new_tokens=64
    ):
        print("\033[H\033[J")
        print("=== 实时流式分析状态 ===")
        print(chunk_output)
        print("-" * 40)