#!/usr/bin/env python3
"""Apply the CoreX vLLM compatibility patches required by the Gemma 4 family.

Both members of the family (Gemma-4-31B-it dense, GEV-26B-Decide MoE) need the
same two fixes, generated against the exact base image

    registry.iluvatar.com.cn:10443/customer/sz/vllm0.17.0-4.4.0-x86:v7
    (vLLM 0.17.0+corex.20260612070235)

1. `vllm/model_executor/models/gemma4_mm.py` imports `MultiModalDataDict`
   from `vllm.inputs`, but this build exposes it from
   `vllm.multimodal.inputs`.  Without the fix, model inspection fails before
   the server starts:
       ImportError: cannot import name 'MultiModalDataDict' from 'vllm.inputs'

2. `vllm/model_executor/layers/fused_moe/fused_moe.py` checks
   `TritonAttentionMetadata.num_decodes` / `num_prefills` for the decode-only
   fast path.  The CoreX Triton attention metadata only has `max_query_len`.
   Gemma-4-31B is dense so this path is normally unused, but GEV (MoE) fails
   CUDA graph capture without the fix:
       AttributeError: 'TritonAttentionMetadata' object has no attribute 'num_decodes'

The script is idempotent, tolerates an already-reformatted target file, and
validates every patched file with `ast.parse`.  It runs both at image build
time and again on container start, so a stale image still gets fixed.
"""

from __future__ import annotations

import ast
import os
import sys
from pathlib import Path

try:
    import vllm
except ImportError as exc:  # pragma: no cover - only in broken images
    print(f"ERROR: cannot import vllm: {exc}", file=sys.stderr)
    raise SystemExit(1)

PKG = Path(
    os.environ.get(
        "VLLM_PKG_ROOT",
        str(Path(vllm.__file__).resolve().parent.parent),
    )
)

MM_FILE = PKG / "vllm/model_executor/models/gemma4_mm.py"
MOE_FILE = PKG / "vllm/model_executor/layers/fused_moe/fused_moe.py"

MM_OLD = "from vllm.inputs import MultiModalDataDict"
MM_NEW = "from vllm.multimodal.inputs import MultiModalDataDict"
MM_MARK = "from vllm.multimodal.inputs import MultiModalDataDict"

MOE_OLD = (
    "        only_decode = (use_ep == False and all("
    "t.num_decodes > 0 and t.num_prefills ==0 for t in "
    "list(attn_metadata.values())))"
)
MOE_NEW = """        only_decode = (
            use_ep == False
            and all(
                getattr(t, "max_query_len", 2) == 1
                for t in list(attn_metadata.values())
            )
        )"""
# 单双引号两种写法都算已打过，避免不同格式化工具造成重复替换。
MOE_MARK = (
    'getattr(t, "max_query_len", 2) == 1',
    "getattr(t, 'max_query_len', 2) == 1",
)


def _patch(
    path: Path,
    old: str,
    new: str,
    marker: str | tuple[str, ...],
    label: str,
) -> None:
    if not path.exists():
        raise SystemExit(f"target file not found: {path}")

    src = path.read_text(encoding="utf-8")
    markers = (marker,) if isinstance(marker, str) else marker
    if any(m in src for m in markers):
        print(f"[skip] {label}: already patched ({path})")
        return
    if old not in src:
        raise SystemExit(
            f"ERROR: {label}: expected code block not found in {path}.\n"
            "The base image or vLLM version may differ; inspect the file manually."
        )

    path.write_text(src.replace(old, new, 1), encoding="utf-8")
    print(f"[ok]   {label}: {path}")


def main() -> None:
    _patch(MM_FILE, MM_OLD, MM_NEW, MM_MARK, "Gemma4 MultiModalDataDict import")
    _patch(MOE_FILE, MOE_OLD, MOE_NEW, MOE_MARK, "TritonAttentionMetadata decode check")

    for path in (MM_FILE, MOE_FILE):
        ast.parse(path.read_text(encoding="utf-8"))
    print("all Gemma 4 family patches applied; syntax OK")


if __name__ == "__main__":
    main()