#!/usr/bin/env python3
"""Target-side EAGLE3/aux-hidden-state support for Glm5Next (DFlash2 needs it).

The DFlash2 drafter consumes the TARGET's hidden states at layers
5/14/24/33/42 (draft config.json dflash_config.target_layer_ids). The vLLM
spec-decode runner gates that on supports_eagle3(target) and on the model
forward emitting aux states (EagleModelMixin). The day-0 nvidia Glm5Next
model implements neither; the fork (local-inference-lab/vllm
dev/jovian-judgement @ da4d7be6, Apache-2.0, see UPSTREAM.md) does. A full
file swap is NOT possible: the fork diverges from our day-0 base in ~150
lines (b12x backends, weight remaps) and our base carries fixes the fork
lacks. So this is a surgical port of exactly the fork's:

  * EagleModelMixin on Glm5NextModel + SupportsEagle3 on both roots
  * dflash_capture flag (fork: speculative_config.use_dflash(), absent on
    our base -> method == "dflash" equivalent)
  * _prepare_aux_hidden_state (mHC: hc_post then hc_contract for dflash)
  * forward-loop aux collection + tuple return (after self.norm, fork parity)

Idempotent; every anchor asserted exactly-once; py_compile-checked.
"""

import py_compile

B = "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia"
PATH = f"{B}/model.py"

EDITS = [
    # 1. imports
    (
        """from vllm.model_executor.models.interfaces import (
    HasInnerState,
    IsHybrid,
    MixtureOfExperts,
    SupportsPP,
)""",
        """from vllm.model_executor.models.interfaces import (
    EagleModelMixin,
    HasInnerState,
    IsHybrid,
    MixtureOfExperts,
    SupportsEagle3,
    SupportsPP,
)""",
    ),
    # 2. model gets the mixin (runner sets aux layers through it)
    (
        "class Glm5NextModel(nn.Module):\n",
        "class Glm5NextModel(nn.Module, EagleModelMixin):\n",
    ),
    # 3. dflash capture flag in Glm5NextModel.__init__
    (
        """        self.is_sequence_parallel = (
            vllm_config.parallel_config.use_sequence_parallel_moe
        )
""",
        """        self.is_sequence_parallel = (
            vllm_config.parallel_config.use_sequence_parallel_moe
        )
        _spec = vllm_config.speculative_config
        # fork parity: dflash drafts consume hc_contract-ed aux states
        # (SpeculativeConfig.use_dflash() only exists on the fork).
        self.dflash_capture = (
            _spec is not None and _spec.method == "dflash"
        )
""",
    ),
    # 4. aux extraction method (fork body, verbatim apart from type-name quote)
    (
        """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)
""",
        """    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

    def _prepare_aux_hidden_state(
        self,
        layer: "Glm5NextDecoderLayer",
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
        post: torch.Tensor | None,
        comb: torch.Tensor | None,
    ) -> torch.Tensor:
        if not layer.mhc or residual is None:
            return hidden_states
        assert post is not None and comb is not None
        aux_hidden_state = layer.hc_post(hidden_states, residual, post, comb)
        if self.dflash_capture:
            return hc_contract(aux_hidden_state, layer.n)
        return aux_hidden_state.flatten(1)
""",
    ),
    # 5. forward return type
    (
        """    ) -> torch.Tensor:
        if get_pp_group().is_first_rank:
""",
        """    ) -> (
        torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]
    ):
        if get_pp_group().is_first_rank:
""",
    ),
    # 6. aux collection in the layer loop (fork semantics: capture the state
    # ENTERING layer_idx+1, i.e. the OUTPUT of layer_idx)
    (
        """        for layer in self._active_layers:
            hidden_states, residual, post, comb = layer(
                positions, hidden_states, residual, post, comb
            )
""",
        """        aux_hidden_states: list[torch.Tensor] = []
        if self.start_layer in self.aux_hidden_state_layers:
            aux_hidden_states.append(hidden_states)

        for layer_idx, layer in enumerate(
            self._active_layers, start=self.start_layer
        ):
            if layer_idx + 1 in self.aux_hidden_state_layers:
                aux_hidden_state = self._prepare_aux_hidden_state(
                    layer, hidden_states, residual, post, comb
                )
                if self.is_sequence_parallel:
                    aux_hidden_state = sp_all_gather(aux_hidden_state)[
                        :full_num_tokens
                    ]
                aux_hidden_states.append(aux_hidden_state)
            hidden_states, residual, post, comb = layer(
                positions, hidden_states, residual, post, comb
            )
""",
    ),
    # 7. tuple return after norm (fork parity — logits stay well-formed)
    (
        """        hidden_states = self.norm(hidden_states)
        return hidden_states

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
""",
        """        hidden_states = self.norm(hidden_states)
        if aux_hidden_states:
            return hidden_states, aux_hidden_states
        return hidden_states

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
""",
    ),
    # 8. both registered roots advertise eagle3 (the runner checks the root)
    (
        """    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
):
""",
        """    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid,
    SupportsEagle3
):
""",
    ),
    (
        """    Glm4vForConditionalGeneration, HasInnerState, IsHybrid
):
""",
        """    Glm4vForConditionalGeneration, HasInnerState, IsHybrid, SupportsEagle3
):
""",
    ),
]

src = open(PATH).read()
if "dflash_capture" in src:
    print(f"already patched: {PATH}")
else:
    for i, (old, new) in enumerate(EDITS, 1):
        n = src.count(old)
        assert n == 1, f"edit {i}: anchor {n}x in {PATH} — review me"
        src = src.replace(old, new, 1)
    open(PATH, "w").write(src)
    py_compile.compile(PATH, doraise=True)
    print(f"patched: {PATH}")

print("patch_eagle3 OK")