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

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)

import torch
from transformers import AutoModelForCausalLM, AutoProcessor


# ===================== 配置 =====================

MODEL_ID = "microsoft/Mage-VL"
VIDEO_FILE = "/content/soccer-broadcast.mp4"

SEGMENT_SEC = 4.0

# Gate 使用：
# 前4秒 + 当前4秒 = 最多8秒视觉上下文
CONTEXT_SEC = 4.0

# LLM 保留最近几条文字事件
HISTORY_SIZE = 3

MAX_SEGMENTS = None          # 测试可改成 8
GATE_THRESHOLD = 0.50
MAX_NEW_TOKENS = 64

CODEC_MAX_PIXELS = 64_000
CODEC_TARGET_CANVAS = 24

SIMULATE_REALTIME = True


# ===================== Prompt =====================

GATE_PROMPT = """
这是足球比赛视频。
请判断视频末尾的当前时间窗口是否包含值得进一步分析的内容。
重点关注传球、带球、射门、抢断、犯规、角球、任意球、进球、
危险进攻、球员明显动作或比赛状态变化。
""".strip()


def make_prompt(start, end, history):
    history_text = "\n".join(
        f"- {x}" for x in history[-HISTORY_SIZE:]
    )

    if not history_text:
        history_text = "暂无历史事件。"

    return f"""
请用中文分析足球比赛视频。

当前重点时间：
{start:.1f}–{end:.1f} 秒

视频包含：
- 前面的视觉内容：用于理解动作上下文
- 当前时间窗口：重点分析当前发生的事情

最近已经观察到的事件：
{history_text}

请重点描述当前时间窗口发生了什么，并结合前面的视觉内容判断动作是否连续。

重点关注：
传球、带球、射门、抢断、犯规、角球、任意球、进球、
危险进攻、球员位置变化以及比赛状态变化。

如果当前只是普通画面，没有明显事件，也请简短说明。

不要猜测无法从视频确认的信息。
尽量简洁，适合实时视频分析。
""".strip()


# ===================== 模型 =====================

if not torch.cuda.is_available():
    raise RuntimeError("未检测到 CUDA")

DEVICE = "cuda"
DTYPE = torch.float16

print(f"[GPU] {torch.cuda.get_device_name(0)}")
print(f"[CUDA] {torch.version.cuda}")

print("[init] 加载 Processor...")

processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

print("[init] 加载 Mage-VL...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=DTYPE,
    attn_implementation="sdpa",
).to(DEVICE).eval()

print("[init] 模型加载完成")

_REMOTE_PKG = type(processor).__module__.rsplit(".", 1)[0]

codec_mod = importlib.import_module(
    _REMOTE_PKG + ".codec_video_processing_mage_vl"
)

CodecConfig = codec_mod.CodecConfig

try:
    model.model._load_streammind_gate()
    print("[gate] Cognition Gate 加载成功")
except Exception as e:
    print(f"[gate] 加载失败: {e}")


# ===================== 工具 =====================

def get_duration(path):
    try:
        r = subprocess.run(
            [
                "ffprobe", "-v", "error",
                "-show_entries", "format=duration",
                "-of", "default=nw=1:nk=1",
                str(path)
            ],
            capture_output=True,
            text=True,
            check=True,
            timeout=60
        )
        return float(r.stdout.strip())
    except Exception:
        return 0.0


def codec_config():
    override = {
        "engine": "hevc",
        "target_canvas": CODEC_TARGET_CANVAS,
        "patch": 16,
    }

    kwargs = dict(processor._codec_config_defaults)
    dcvc = dict(kwargs.get("dcvc") or {})

    kwargs.update(override)

    if dcvc:
        kwargs["dcvc"] = dcvc

    kwargs["max_pixels"] = CODEC_MAX_PIXELS

    return CodecConfig(**kwargs), override


def make_chat_prompt(text):
    messages = [{
        "role": "user",
        "content": [
            {"type": "video"},
            {"type": "text", "text": text},
        ],
    }]

    return processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )


def extract_clip(src, start, length, out):
    tmp = out.with_name(".tmp_" + out.name)

    if tmp.exists():
        tmp.unlink()

    subprocess.run(
        [
            "ffmpeg", "-y", "-loglevel", "error",
            "-ss", f"{max(0, start):.3f}",
            "-i", str(src),
            "-t", f"{length:.3f}",
            "-c:v", "libx264",
            "-preset", "ultrafast",
            "-crf", "23",
            "-an",
            "-pix_fmt", "yuv420p",
            str(tmp)
        ],
        check=True,
        capture_output=True,
        timeout=180
    )

    tmp.replace(out)
    return out


def to_cuda(inputs):
    result = {}

    for k, v in inputs.items():
        if not hasattr(v, "to"):
            continue

        v = v.to(DEVICE)

        if k == "pixel_values":
            v = v.to(DTYPE)

        result[k] = v

    return result


def cleanup(*objs):
    for obj in objs:
        try:
            del obj
        except Exception:
            pass

    gc.collect()
    torch.cuda.empty_cache()


def gpu_mem():
    a = torch.cuda.memory_allocated() / 1024**3
    r = torch.cuda.memory_reserved() / 1024**3
    return f"{a:.2f}GB / reserved {r:.2f}GB"


# ===================== Gate =====================

@torch.inference_mode()
def run_gate(inputs):
    feat = {}

    for k in (
        "pixel_values",
        "image_grid_thw",
        "patch_positions",
    ):
        if k not in inputs:
            continue

        v = inputs[k].to(DEVICE)

        if k == "pixel_values":
            v = v.to(DTYPE)

        feat[k] = v

    logits = None

    try:
        logits = model.model.streammind_gate_forward_segments(
            [feat]
        )[0]

        # 这里取视觉序列最后位置。
        #
        # 因为输入是：
        # 前4秒 + 当前4秒
        #
        # 所以最后位置对应当前4秒末尾，
        # Gate 可以利用前面的视觉内容作为上下文。

        length = int(
            inputs["image_grid_thw"][:, 0].sum()
        )

        prob = torch.softmax(
            logits[length - 1].float(),
            dim=-1
        )[1].item()

        return float(prob)

    finally:
        cleanup(feat, logits)


# ===================== Codec =====================

def encode_video(
    clip,
    prompt,
    codec_override
):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        return processor(
            text=[prompt],
            videos=[str(clip)],
            video_backend="codec",
            max_pixels=CODEC_MAX_PIXELS,
            codec_config=codec_override,
            return_tensors="pt",
            padding=False,
        )


# ===================== Mage-VL =====================

@torch.inference_mode()
def generate(inputs, max_new_tokens):
    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        use_cache=True,
    )

    n = inputs["input_ids"].shape[1]

    text = processor.tokenizer.decode(
        out[0, n:],
        skip_special_tokens=True
    ).strip()

    del out

    return text


# ===================== 主循环 =====================

def stream_video_analysis(
    video_path,
    segment_sec=SEGMENT_SEC,
    context_sec=CONTEXT_SEC,
    gate_threshold=GATE_THRESHOLD,
    max_segments=MAX_SEGMENTS,
    max_new_tokens=MAX_NEW_TOKENS,
    simulate_realtime=SIMULATE_REALTIME,
):
    path = Path(video_path)

    if not path.exists():
        raise FileNotFoundError(video_path)

    total = get_duration(path)

    if total <= 0:
        raise RuntimeError("无法获取视频时长")

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

    _, codec_override = codec_config()

    start = 0.0
    index = 0

    # 最近几次文字事件
    history = []

    # 所有输出
    results = []

    yield (
        f"开始实时流式分析\n"
        f"视频={total:.1f}s | "
        f"当前窗口={segment_sec:.1f}s | "
        f"视觉上下文={context_sec:.1f}s | "
        f"历史={HISTORY_SIZE}条"
    )

    while start < total - 0.05:

        if (
            max_segments is not None
            and index >= max_segments
        ):
            break

        # 当前时间窗口
        end = min(
            total,
            start + segment_sec
        )

        current_length = end - start

        # --------------------------------------------------
        # 视觉上下文：
        #
        # 第1段：
        # 0-4
        #
        # 第2段：
        # 0-8
        #
        # 第3段：
        # 4-12
        #
        # 第4段：
        # 8-16
        # --------------------------------------------------

        context_start = max(
            0.0,
            start - context_sec
        )

        context_end = end

        visual_length = (
            context_end - context_start
        )

        index += 1

        t0 = time.time()

        clip = clip_dir / (
            f"{path.stem}_"
            f"{index:06d}_"
            f"{context_start:.2f}_"
            f"{context_end:.2f}.mp4"
        )

        inputs = None
        cuda_inputs = None

        try:

            # ==============================================
            # 1. 当前滑动视觉窗口
            # ==============================================

            extract_clip(
                path,
                context_start,
                visual_length,
                clip
            )

            # ==============================================
            # 2. Gate
            #
            # Gate 看到：
            #
            # 前4秒 + 当前4秒
            # ==============================================

            gate_prompt = make_chat_prompt(
                GATE_PROMPT
            )

            inputs = encode_video(
                clip,
                gate_prompt,
                codec_override
            )

            prob = run_gate(inputs)

            head = (
                f"**{start:.1f}–{end:.1f}s** "
                f"Gate=`{prob:.2f}` "
                f"[context={context_start:.1f}–{end:.1f}s]"
            )

            # ==============================================
            # 3. Gate 不通过
            # ==============================================

            if prob < gate_threshold:

                elapsed = time.time() - t0
                speed = current_length / elapsed

                line = (
                    f"{head} → 🔇 跳过 "
                    f"({elapsed:.2f}s / "
                    f"{speed:.2f}×)"
                )

                results.append(line)

                yield "\n\n".join(results)

            # ==============================================
            # 4. Gate 通过
            # ==============================================

            else:

                yield (
                    "\n\n".join(results)
                    + (
                        "\n\n"
                        if results else ""
                    )
                    + f"{head} → 🧠 生成中..."
                )

                # ------------------------------------------
                # LLM 使用：
                #
                # 视觉：
                # 前4秒 + 当前4秒
                #
                # 文本：
                # 最近3个事件
                # ------------------------------------------

                llm_prompt = make_prompt(
                    start,
                    end,
                    history
                )

                # 重新生成带历史 prompt 的输入
                #
                # 这里重新 processor 是为了让：
                # Gate prompt 和 LLM prompt 分开。
                #
                # 视频 Codec 本身仍然是同一个滑动窗口。
                # ------------------------------------------

                cleanup(inputs)
                inputs = None

                inputs = encode_video(
                    clip,
                    make_chat_prompt(llm_prompt),
                    codec_override
                )

                gt = time.time()

                cuda_inputs = to_cuda(
                    inputs
                )

                said = generate(
                    cuda_inputs,
                    max_new_tokens
                )

                gen_time = time.time() - gt
                elapsed = time.time() - t0
                speed = current_length / elapsed

                # ------------------------------------------
                # 保存历史事件
                # ------------------------------------------

                history.append(
                    f"{start:.1f}–{end:.1f}s：{said}"
                )

                if len(history) > HISTORY_SIZE:
                    history.pop(0)

                line = (
                    f"{head} → 🗣️ {said}\n"
                    f"   Generate={gen_time:.2f}s | "
                    f"Total={elapsed:.2f}s | "
                    f"Realtime={speed:.2f}×"
                )

                results.append(line)

                yield "\n\n".join(results)

            # ==============================================
            # 5. 模拟实时
            # ==============================================

            elapsed = time.time() - t0

            if (
                simulate_realtime
                and elapsed < current_length
            ):
                time.sleep(
                    current_length - elapsed
                )

        except Exception as e:

            results.append(
                f"**{start:.1f}–{end:.1f}s** "
                f"→ ❌ "
                f"{type(e).__name__}: {e}"
            )

            yield "\n\n".join(results)

        finally:

            cleanup(
                cuda_inputs,
                inputs
            )

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

            start = end

            print(
                f"[segment {index}] "
                f"{start:.1f}/{total:.1f}s | "
                f"GPU {gpu_mem()}",
                flush=True
            )

    yield (
        "\n\n".join(results)
        + f"\n\n✅ 完成，共 {index} 个片段"
    )


# ===================== Main =====================

if __name__ == "__main__":

    for output in stream_video_analysis(
        VIDEO_FILE,
        segment_sec=4.0,
        context_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 第六版：上下文实时分析 ===\n")
        print(output)
        print("\n" + "-" * 70)