#!/usr/bin/env python3
"""Apply the CoreX vLLM patches needed by Llama 4 Scout / Maverick.

The patches in ./patches were 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).

Scout needs only the chunked-attention patch.  Maverick additionally needs the
FP8 quantization / FP8 MoE patches and the `corexx_fp8_moe` fallback module.
One image contains the union of both patch sets; the FP8 paths are only used
when the mounted model is FP8.

The script is idempotent: it checks for a marker in each target file and skips
patches that are already applied.
"""

from __future__ import annotations

import ast
import os
import pathlib
import shutil
import subprocess
import sys

PATCH_DIR = pathlib.Path(
    os.environ.get("LLAMA4_PATCH_DIR", "/opt/llama4/patches")
)

# LLAMA4_PKG_ROOT is only needed for local testing against a copy of the
# vendor files.  In the image it is derived from the installed vllm package.
_pkg_env = os.environ.get("LLAMA4_PKG_ROOT")
if _pkg_env:
    PKG = pathlib.Path(_pkg_env)
else:
    import vllm

    PKG = pathlib.Path(vllm.__file__).resolve().parent.parent

# patch file -> (target file relative to the package root, applied marker)
TARGETS = {
    "01_fp8_quant.patch": (
        "vllm/_custom_ops.py",
        "COREXX_FP8_QUANT_PATCH",
    ),
    "02_fp8_moe.patch": (
        "vllm/model_executor/layers/fused_moe/fused_moe.py",
        "COREXX_FP8_MOE_PATCH",
    ),
    "03_llama4_chunked_attn.patch": (
        "vllm/v1/attention/backends/utils.py",
        "seq_lens_np=seqlens_k_local",
    ),
}


def apply_one(patch_name: str, target_rel: str, marker: str) -> None:
    target = PKG / target_rel
    if not target.exists():
        sys.exit(f"target file not found: {target}")
    if marker in target.read_text():
        print(f"[skip] {target} (already patched)")
        return

    patch_file = PATCH_DIR / patch_name
    if not patch_file.exists():
        sys.exit(f"missing patch file: {patch_file}")

    result = subprocess.run(
        [
            "patch",
            "-p1",
            "--forward",
            "--batch",
            "--no-backup-if-mismatch",
            "-d",
            str(PKG),
            "-i",
            str(patch_file),
        ],
        text=True,
        capture_output=True,
    )
    if result.stdout:
        print(result.stdout, end="")
    if result.returncode != 0:
        if result.stderr:
            print(result.stderr, file=sys.stderr)
        sys.exit(f"patch failed: {patch_name}")
    if marker not in target.read_text():
        sys.exit(f"patch {patch_name} applied but marker not found in {target}")
    print(f"[ok]   {target}")


def main() -> None:
    for patch_name, (target_rel, marker) in TARGETS.items():
        apply_one(patch_name, target_rel, marker)

    src = PATCH_DIR / "corexx_fp8_moe.py"
    dst = PKG / "vllm/model_executor/layers/fused_moe/corexx_fp8_moe.py"
    if not src.exists():
        sys.exit(f"missing module: {src}")
    shutil.copy2(src, dst)
    print(f"[copy] {src} -> {dst}")

    for path in (
        PKG / "vllm/_custom_ops.py",
        PKG / "vllm/model_executor/layers/fused_moe/fused_moe.py",
        dst,
        PKG / "vllm/v1/attention/backends/utils.py",
    ):
        ast.parse(path.read_text())
    print("all patches applied; syntax OK")


if __name__ == "__main__":
    main()