#!/usr/bin/env python3
"""Enable EAGLE-3 for the Qwen3.5/3.6 target (Qwen3_5ForConditionalGeneration) in SGLang.

Run once:
    python3 patch_sglang_eagle3.py            # auto-locates installed sglang
    python3 patch_sglang_eagle3.py <path-to qwen3_5.py>

Idempotent, writes a .eagle3-bak backup, AST-validates the result.
Requires SGLang >= 0.5.10.
"""
import ast
import importlib.util
import os
import shutil
import sys

MARKER = "[EAGLE3-PATCH]"

DECODER_ANCHOR = (
    "    def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):\n"
    "        self.layers_to_capture = layers_to_capture\n"
    "        for layer_id in self.layers_to_capture:\n"
    '            setattr(self.layers[layer_id], "_is_layer_to_capture", True)\n'
)
DECODER_NEW = DECODER_ANCHOR + (
    "\n"
    "    def set_eagle3_layers_to_capture(self, layers_to_capture: list[int]):\n"
    f"        # {MARKER}\n"
    "        self.layers_to_capture = layers_to_capture\n"
    "        for layer_id in self.layers_to_capture:\n"
    '            setattr(self.layers[layer_id], "_is_layer_to_capture", True)\n'
)

WRAPPER_ANCHOR = (
    "class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):"
)
WRAPPER_NEW = (
    "    def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None):\n"
    f"        # {MARKER}\n"
    "        self.capture_aux_hidden_states = True\n"
    "        if layer_ids is None:\n"
    '            text_cfg = getattr(self.config, "text_config", self.config)\n'
    "            num_layers = text_cfg.num_hidden_layers\n"
    "            offset_ids = [2, num_layers // 2, num_layers - 3]\n"
    "        else:\n"
    "            offset_ids = [val + 1 for val in layer_ids]\n"
    "        self.model.set_eagle3_layers_to_capture(offset_ids)\n"
    "\n"
    "\n"
    + WRAPPER_ANCHOR
)


def locate_qwen3_5():
    spec = importlib.util.find_spec("sglang")
    if spec is None or not getattr(spec, "submodule_search_locations", None):
        sys.exit("ERROR: 'sglang' not found. Activate the env or pass the qwen3_5.py path.")
    path = os.path.join(list(spec.submodule_search_locations)[0], "srt", "models", "qwen3_5.py")
    if not os.path.isfile(path):
        sys.exit(f"ERROR: {path} not found -- SGLang >= 0.5.10 required.")
    return path


def main():
    if len(sys.argv) == 1:
        path = locate_qwen3_5()
    elif len(sys.argv) == 2:
        path = sys.argv[1]
    else:
        sys.exit("usage: patch_sglang_eagle3.py [path-to qwen3_5.py]")

    with open(path) as f:
        src = f.read()

    if MARKER in src:
        print(f"{path}: already patched.")
        return

    if src.count(DECODER_ANCHOR) != 1:
        sys.exit("FAIL: decoder anchor not found (unexpected SGLang version).")
    if src.count(WRAPPER_ANCHOR) != 1:
        sys.exit("FAIL: wrapper anchor not found (unexpected SGLang version).")

    patched = src.replace(DECODER_ANCHOR, DECODER_NEW, 1).replace(WRAPPER_ANCHOR, WRAPPER_NEW, 1)
    try:
        ast.parse(patched)
    except SyntaxError as e:
        sys.exit(f"FAIL: patched file does not parse: {e}")
    if patched.count("def set_eagle3_layers_to_capture") != 2:
        sys.exit("FAIL: expected 2 injected methods.")

    bak = path + ".eagle3-bak"
    if not os.path.exists(bak):
        shutil.copyfile(path, bak)
    with open(path, "w") as f:
        f.write(patched)
    print(f"{path}: patched OK")


if __name__ == "__main__":
    main()