"""Teach SGLang-Omni's Drama 3 loader the full int4 checkpoint (quant_method "drama3-int4").

Routed FusedMoE experts run SGLang's moe_wna16 kernel (W4A16, GPTQ packing, symmetric, group 128);
the dense target linears (full-attention q/k/v/o, gated-deltanet in_proj_qkv/in_proj_z/out_proj,
shared-expert gate/up/down) run GPTQ-Marlin (standard GPTQ sym checkpoint fields, repacked by
GPTQMarlinLinearMethod.process_weights_after_loading); every other layer stays unquantized bf16.

Idempotent. Usage: python apply_patch.py /path/to/sglang-omni
"""
import sys
from pathlib import Path

root = Path(sys.argv[1]) / "sglang_omni/models/fishaudio_drama3"
q = root / "quantization.py"
s = q.read_text()
ADD = '''

# --- added by drama-3-int4 (int4 experts AND dense linears) ----------------------------------
import re as _re  # noqa: E402

from sglang.srt.layers.quantization.gptq.gptq import (  # noqa: E402
    GPTQMarlinConfig,
    GPTQMarlinLinearMethod,
)
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config, MoeWNA16Method  # noqa: E402

DRAMA3_INT4_QUANTIZATION_METHOD = "drama3-int4"

# Dense checkpoint-side prefixes that hold int4 GPTQ fields. Matches by checkpoint name, which
# is what sglang's weight_loader prefix strings carry (model.language_model.layers.N....).
# The checkpoint's quantization_config.dense.targets says exactly which dense linears were
# quantized offline; the default below matches the full drama3-int4 release. Never matches
# mlp.gate (router), mlp.shared_expert_gate, linear_attn.in_proj_a/b, conv1d, norms, embeddings.
_ALL_DENSE_TARGETS = [
    "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
    "linear_attn.in_proj_qkv", "linear_attn.in_proj_z", "linear_attn.out_proj",
    "mlp.shared_expert.gate_proj", "mlp.shared_expert.up_proj", "mlp.shared_expert.down_proj",
]


class Drama3Int4Config(QuantizationConfig):
    """int4 W4A16 everywhere it pays: moe_wna16 for FusedMoE, GPTQ-Marlin for dense targets."""

    # SGLang fuses modules at load time; the loader prefix is the FUSED name while the
    # checkpoint (and our targets list) uses per-projection names. Map back before matching.
    MERGED = {"qkv_proj": ("q_proj", "k_proj", "v_proj"),
              "gate_up_proj": ("gate_proj", "up_proj"),
              "in_proj_qkvz": ("in_proj_qkv", "in_proj_z")}

    def _checkpoint_names(self, prefix: str) -> list:
        p = prefix.replace("model.layers.", "model.language_model.layers.", 1)
        head, _, last = p.rpartition(".")
        return [f"{head}.{n}" for n in self.MERGED.get(last, (last,))]

    def __init__(self, group_size: int = 128, targets: list | None = None) -> None:
        super().__init__()
        full_moe = {"quant_method": "gptq", "bits": 4, "group_size": group_size,
                    "sym": True, "desc_act": False}
        self.wna16 = MoeWNA16Config("gptq", 4, group_size, False, False, None, full_moe)
        full_dense = {"quant_method": "gptq_marlin", "bits": 4, "group_size": group_size,
                      "sym": True, "desc_act": False}
        self.marlin = GPTQMarlinConfig(4, group_size, False, True, False, {}, full_dense)
        targets = _ALL_DENSE_TARGETS if targets is None else targets
        self.dense_re = (_re.compile("|".join(_re.escape(t) for t in targets))
                         if targets else None)

    def get_name(self) -> str:
        return DRAMA3_INT4_QUANTIZATION_METHOD

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.bfloat16, torch.half]

    @classmethod
    def get_min_capability(cls) -> int:
        return 80

    @staticmethod
    def get_config_filenames() -> list[str]:
        return []

    @classmethod
    def from_config(cls, config: Mapping[str, object]) -> "Drama3Int4Config":
        if (config.get("quant_method") != DRAMA3_INT4_QUANTIZATION_METHOD
                or config.get("bits") != 4 or not config.get("sym")):
            raise ValueError(
                f"expected quant_method={DRAMA3_INT4_QUANTIZATION_METHOD!r}, bits=4, sym=True; "
                f"got {dict(config)}"
            )
        gs = int(config.get("group_size", 128))
        if gs <= 0 or gs % 32:
            raise ValueError(f"group_size must be a positive multiple of 32; got {gs}")
        dense_cfg = config.get("dense") or {}
        targets = dense_cfg.get("targets") if isinstance(dense_cfg, Mapping) else None
        return cls(gs, targets)

    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
        if isinstance(layer, FusedMoE):
            return MoeWNA16Method(self.wna16)
        elif isinstance(layer, LinearBase):
            if (prefix and self.dense_re
                    and any(self.dense_re.search(n) for n in self._checkpoint_names(prefix))):
                return GPTQMarlinLinearMethod(self.marlin)
            return UnquantizedLinearMethod()
        else:
            return None

    def get_scaled_act_names(self) -> list[str]:
        return []
'''
if "Drama3Int4Config" not in s:
    q.write_text(s + ADD)
e = root / "engine_builder.py"
s = e.read_text()
if "Drama3Int4Config" not in s:
    s = s.replace("    ExpertFp8Config,\n", "    ExpertFp8Config,\n    DRAMA3_INT4_QUANTIZATION_METHOD,\n    Drama3Int4Config,\n", 1)
    s = s.replace("        QUANTIZATION_METHODS[QUANTIZATION_METHOD] = ExpertFp8Config\n",
                  "        QUANTIZATION_METHODS[QUANTIZATION_METHOD] = ExpertFp8Config\n"
                  "        QUANTIZATION_METHODS[DRAMA3_INT4_QUANTIZATION_METHOD] = Drama3Int4Config\n", 1)
    assert "QUANTIZATION_METHODS[DRAMA3_INT4_QUANTIZATION_METHOD]" in s, "engine_builder layout changed"
    e.write_text(s)
print("patched", q, e)