#!/usr/bin/env python3
"""Patch vLLM's set_eagle3_aux_hidden_state_layers to resolve the eagle3
interface on the inner language-model decoder when the outer model is a
wrapper (e.g. MiMoV2OmniForCausalLM, whose text stack is
MiMoV2ForCausalLM/MiMoV2Model and does implement SupportsEagle3).

Root cause (2026-09-28, mimo26 vLLM lane): MiMo-V2.6-Flash-MOPD resolves to
MiMoV2OmniForCausalLM (multimodal wrapper). For method=dflash the GPU model
runner marks use_aux_hidden_state_outputs=True and calls
set_eagle3_aux_hidden_state_layers(target, spec), which isinstance-checks
SupportsEagle3 against the OUTER class and raises
"Model does not support EAGLE3 interface" before ever consulting
_inner_decoder(). Everything downstream (aux slots, _maybe_add_hidden_state)
lives on the inner MiMoV2Model, and with PP=1 the relay helpers are no-ops,
so an inner-resolved check is exactly the correct behaviour here.

Idempotent, anchor-based: `--check` exits 0 only when the patch is present,
2 when the file is still patchable (and the build must apply it), 1 when the
anchor is absent — it refuses silently-wrong sources. `--apply` writes once
(no backup file; the container layer is disposable and the image rebuilds
from this dir).
"""

import sys

MODULE = "vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils"

ANCHOR = (
    "    if not supports_eagle3(model):\n"
    '        raise RuntimeError("Model does not support EAGLE3 interface")\n'
)

PATCHED = (
    "    if not supports_eagle3(model):\n"
    "        # Multimodal wrappers (e.g. MiMoV2OmniForCausalLM) keep the\n"
    "        # eagle3 carrier one level in — the language model whose method\n"
    "        # set_aux_hidden_state_layers() already knows how to reach the\n"
    "        # EagleModelMixin holder. Resolve it before failing (the\n"
    "        # _inner_decoder() helper does NOT fit: it returns the bare\n"
    "        # decoder, which lacks the protocol's method).\n"
    "        _lang = None\n"
    "        if hasattr(model, \"get_language_model\"):\n"
    "            _lang = model.get_language_model()\n"
    "        elif hasattr(model, \"language_model\"):\n"
    "            _lang = model.language_model\n"
    "        if _lang is None or not supports_eagle3(_lang):\n"
    '            raise RuntimeError("Model does not support EAGLE3 interface")\n'
    "        logger.info(\n"
    '            "Eagle3 interface resolved on %s inside %s",\n'
    "            type(_lang).__name__,\n"
    "            type(model).__name__,\n"
    "        )\n"
    "        model = _lang\n"
)


def target_path() -> str:
    import importlib

    mod = importlib.import_module(MODULE)
    path = getattr(mod, "__file__", None)
    if not path:
        raise SystemExit(f"{MODULE} has no __file__ (frozen install?)")
    return path


def main() -> int:
    check = "--check" in sys.argv
    path = target_path()
    with open(path, encoding="utf-8") as f:
        src = f.read()

    if PATCHED in src:
        print(f"ALREADY PATCHED: {path}")
        return 0
    if ANCHOR not in src:
        print(f"ANCHOR NOT FOUND (unexpected source): {path}", file=sys.stderr)
        return 1
    if check:
        print(f"NOT PATCHED YET: {path}", file=sys.stderr)
        return 2

    with open(path, "w", encoding="utf-8") as f:
        f.write(src.replace(ANCHOR, PATCHED, 1))
    print(f"APPLIED: {path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())