#!/usr/bin/env python3
"""
Usage:
    podman run --rm -v $PWD:/o --entrypoint bash <image> -c \
      'cp /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/qwen3_5.py /o/qwen3_5.py.orig; \
       cp /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/qwen3_5_mtp.py /o/qwen3_5_mtp.py.orig'
    python3 make_embed_quant_patch.py
    podman run ... -v $PWD/qwen3_5.py:/usr/local/.../models/qwen3_5.py:ro ...
"""
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))

EDITS = [
    (
        "qwen3_5.py",
        """        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
        )
""",
        """        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            quant_config=self.quant_config,
            prefix=maybe_prefix(prefix, "embed_tokens"),
        )
""",
    ),
    (
        "qwen3_5_mtp.py",
        """        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
        )
""",
        """        self.embed_tokens = VocabParallelEmbedding(
            self.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "embed_tokens"),
        )
""",
    ),
]


def main():
    rc = 0
    for fn, old, new in EDITS:
        src = os.path.join(HERE, fn + ".orig")
        dst = os.path.join(HERE, fn)
        if not os.path.exists(src):
            print(f"✗ {src} が無い。先にイメージから取り出すこと")
            rc = 1
            continue
        text = open(src).read()
        n = text.count(old)
        if n != 1:
            print(f"✗ {fn}: 想定パターンが {n} 個。vLLM 側が変わった可能性がある")
            rc = 1
            continue
        if "maybe_prefix" not in text:
            print(f"✗ {fn}: maybe_prefix が import されていない")
            rc = 1
            continue
        open(dst, "w").write(text.replace(old, new))
        print(f"✓ {fn}: embed_tokens に quant_config / prefix を配線した -> {dst}")
    return rc


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