"""Give Qwen3.8 Flash-Next output heads their configured quantization method."""

from __future__ import annotations

import argparse
import re
from pathlib import Path


MARKER = "qwen38-flash-dgx: support a quantized lm_head checkpoint"
CONSTRUCTOR_HOOK = re.compile(
    r'^(?P<indent>[ \t]+)config\.hidden_size,\n'
    r'(?P=indent)prefix=maybe_prefix\(prefix, "lm_head"\),$',
    re.MULTILINE,
)


def patch_source(source: str) -> str:
    if MARKER in source:
        return source

    def replacement(match: re.Match) -> str:
        indent = match.group("indent")
        return (
            f"{indent}config.hidden_size,\n"
            f"{indent}# {MARKER}\n"
            f"{indent}quant_config=self.quant_config,\n"
            f'{indent}prefix="lm_head",'
        )

    patched, occurrences = CONSTRUCTOR_HOOK.subn(replacement, source)
    if occurrences != 1:
        raise RuntimeError(
            "expected exactly one Qwen output-head constructor hook, "
            f"found {occurrences}"
        )
    return patched


def main(argv=None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("model_paths", type=Path, nargs="+")
    args = parser.parse_args(argv)
    for path in args.model_paths:
        source = path.read_text(encoding="utf-8")
        patched = patch_source(source)
        if patched != source:
            path.write_text(patched, encoding="utf-8")
            print(f"patched quantized Qwen output head: {path}")
        else:
            print(f"quantized Qwen output-head support already present: {path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())