#!/usr/bin/env python3
"""Apply the Agnes patch to an sglang python package directory (the one that
contains `srt/`).  Idempotent.

  1. srt/configs/agnes.py                    new file
  2. srt/utils/hf_transformers/common.py     register AgnesConfig in _CONFIG_REGISTRY
  3. srt/models/qwen3_5.py                   load_weights: translate the Agnes checkpoint
                                              (tensor prefixes, parallel-FFN fold)

Usage: apply_patch.py <path/to/sglang>            e.g. .../site-packages/sglang
"""

import os
import shutil
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
MARK = "# === agnes ==="

TRANSLATE = '''
# === agnes ===
# Agnes 3.0 Flash checkpoints (config model_type "agnes") use their own tensor
# prefixes and carry a parallel FFN branch per layer.  This generator sits at
# the top of the weight stream and turns it into what the implementation below
# expects: delta_attn -> linear_attn, global_attn -> self_attn, and the branch
# concatenated onto the main gate / up (dim 0) and down (dim 1) projections,
# matching the widened intermediate_size set by sglang.srt.configs.agnes.
import json as _agnes_json
import os as _agnes_os
import re as _agnes_re

_AGNES_MLP_RE = _agnes_re.compile(r"^(.*\\.layers\\.\\d+\\.mlp\\.)(gate_proj|up_proj|down_proj)\\.weight$")


class _AgnesBranchReader:
    def __init__(self, model_path):
        from safetensors import safe_open

        self._open = safe_open
        self.path = model_path
        index = _agnes_os.path.join(model_path, "model.safetensors.index.json")
        self.weight_map = _agnes_json.load(open(index))["weight_map"]
        self.handles = {}

    def get(self, key):
        fn = self.weight_map[key]
        if fn not in self.handles:
            self.handles[fn] = self._open(_agnes_os.path.join(self.path, fn), framework="pt", device="cpu")
        return self.handles[fn].get_tensor(key)


def _agnes_translate_weights(model, weights):
    cfg = getattr(model.config, "text_config", None) or model.config
    width = int(getattr(cfg, "agnes_parallel_ffn_intermediate_size", 0) or 0)
    if width <= 0:
        yield from weights
        return
    model_path = (
        getattr(cfg, "agnes_model_path", None)
        or getattr(model.config, "agnes_model_path", None)
        or _agnes_os.environ.get("AGNES_MODEL_PATH")
        or getattr(model.config, "_name_or_path", None)
    )
    if not model_path or not _agnes_os.path.isdir(model_path):
        raise RuntimeError(
            f"agnes: cannot locate the checkpoint directory (got {model_path!r}); "
            "set AGNES_MODEL_PATH to the model directory"
        )
    reader = _AgnesBranchReader(model_path)
    for name, w in weights:
        if ".mlp.parallel_ffn." in name:
            continue
        m = _AGNES_MLP_RE.match(name)
        if m and "visual" not in name and not name.startswith("mtp"):
            extra = reader.get(f"{m.group(1)}parallel_ffn.{m.group(2)}.weight")
            dim = 1 if m.group(2) == "down_proj" else 0
            w = torch.cat([w, extra.to(device=w.device, dtype=w.dtype)], dim=dim)
        name = name.replace(".delta_attn.", ".linear_attn.").replace(".global_attn.", ".self_attn.")
        yield name, w
# === /agnes ===

'''

REGISTER = '''
# === agnes ===
from sglang.srt.configs.agnes import AgnesConfig as _AgnesConfig

_CONFIG_REGISTRY[_AgnesConfig.model_type] = _AgnesConfig
'''


def patch_file(path, edit):
    src = open(path, encoding="utf-8").read()
    if MARK in src:
        return "already patched"
    out = edit(src)
    if out is None:
        raise SystemExit(f"anchor not found in {path}")
    open(path, "w", encoding="utf-8").write(out)
    return "patched"


def edit_model_file(src):
    anchor = "QWEN3_5_KV_SCALE_MAPPER = WeightsMapper("
    hook = "        weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)\n"
    if anchor not in src or src.count(hook) < 1:
        return None
    src = src.replace(anchor, TRANSLATE + anchor, 1)
    src = src.replace(hook, "        weights = _agnes_translate_weights(self, weights)\n" + hook)
    return src


def main():
    if len(sys.argv) != 2:
        sys.exit(__doc__)
    pkg = os.path.abspath(sys.argv[1])
    srt = os.path.join(pkg, "srt")
    if not os.path.isdir(srt):
        sys.exit(f"{pkg} does not contain srt/")
    dst = os.path.join(srt, "configs", "agnes.py")
    shutil.copy2(os.path.join(HERE, "agnes_sglang_config.py"), dst)
    print(f"configs/agnes.py: installed")
    print("utils/hf_transformers/common.py:", patch_file(
        os.path.join(srt, "utils", "hf_transformers", "common.py"), lambda s: s.rstrip("\n") + "\n" + REGISTER))
    print("models/qwen3_5.py:", patch_file(os.path.join(srt, "models", "qwen3_5.py"), edit_model_file))
    print("APPLY_PATCH_OK")


if __name__ == "__main__":
    main()