"""
Patch a modelopt-exported llm-jp-4 checkpoint so it loads/serves correctly
: modelopt's exporter only writes weights + quant config; the
custom tokenizer plugins and chat template do not survive it.

Steps:
  1. Copy auxiliary files from the source model dir into the export dir
     (llmjp4_tokenizer.py, llmjp4_harmony.py, tokenizer files incl.
     tokenizer_config.json whose auto_map wires up Llmjp4Tokenizer,
     chat_template.jinja, generation_config.json).
  2. Verify config.json still declares the right architecture and carries a
     quantization config.
  3. Weight-scale integrity gate: scan all FP8-E4M3 *weight_scale* tensors
     for NaN byte encodings (scan_weight_scale_nan.py); hard-fail if any.
  4. Tokenization round-trip test: encode+decode 50 Japanese strings with the
     PATCHED export dir's tokenizer (trust_remote_code) and require identity
     with the source tokenizer's output.
  5. Print the vLLM serve command for the final smoke test (transformers
     cannot execute NVFP4 checkpoints natively; serving IS the load test).

Usage:
  python3 04_export_and_patch.py --export-dir <export-dir>
"""

import argparse
import json
import shutil
from pathlib import Path

DEFAULT_SOURCE = "llm-jp/llm-jp-4-8b-instruct"
DEFAULT_REVISION = "098f2b2cf33021eba19a6d3582aa3d071ccc0aff"  # the revision this checkpoint was built from

AUX_FILES = [
    "llmjp4_tokenizer.py",
    "llmjp4_harmony.py",
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "chat_template.jinja",
    "generation_config.json",
]

ROUND_TRIP_STRINGS = [
    "こんにちは、世界！",
    "日本語のトークナイザーが正しく動作するか確認します。",
    "御社の益々のご発展をお祈り申し上げます。",
    "量子化されたモデルでも敬語は崩れないはずです。",
    "<|start|>assistant<|channel|>final<|message|>テスト<|return|>",
    "改行\nとタブ\tと　全角スペース",
    "숫자123と英語mixedテキストand絵文字🎌",
    "「鬼滅の刃」の映画は2020年に公開された。",
] + [f"ケース{i}：これは第{i}番目の検証文です。円周率は3.14159…" for i in range(42)]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--export-dir", required=True)
    parser.add_argument("--source-model", default=DEFAULT_SOURCE,
                        help="HF id or local path of the BF16 base model")
    parser.add_argument("--revision", default=DEFAULT_REVISION,
                        help="HF revision of the base model (ignored for "
                             "local paths)")
    args = parser.parse_args()
    from huggingface_hub import snapshot_download

    source = Path(args.source_model) if Path(args.source_model).is_dir() else Path(
        snapshot_download(args.source_model, revision=args.revision)
    )
    export_dir = Path(args.export_dir)
    if not export_dir.is_dir():
        raise SystemExit(f"{export_dir} does not exist")

    # 1. Copy aux files
    for name in AUX_FILES:
        src = source / name
        if not src.exists():
            print(f"  WARNING: source file missing, skipped: {name}")
            continue
        shutil.copy2(src, export_dir / name)
        print(f"  copied {name}")

    # 2. Config checks
    cfg = json.loads((export_dir / "config.json").read_text())
    if cfg.get("architectures") != ["LlamaForCausalLM"]:
        raise SystemExit(f"unexpected architectures: {cfg.get('architectures')}")
    has_quant = "quantization_config" in cfg or (export_dir / "hf_quant_config.json").exists()
    if not has_quant:
        raise SystemExit("no quantization config found in export — modelopt export incomplete?")
    tok_cfg = json.loads((export_dir / "tokenizer_config.json").read_text())
    auto_map = tok_cfg.get("auto_map", {}).get("AutoTokenizer")
    if not (auto_map and "llmjp4_tokenizer" in str(auto_map)):
        raise SystemExit(f"tokenizer auto_map broken: {auto_map}")
    print("config.json + tokenizer_config.json checks passed")

    # 3. Weight-scale integrity gate: modelopt ≤0.44 could emit E4M3 NaN bytes
    # (0x7F/0xFF) when a per-block scale rounds above the FP8 max of 448; one
    # NaN in any weight_scale collapses served output.
    import subprocess
    import sys

    scan = subprocess.run(
        [sys.executable, str(Path(__file__).parent / "scan_weight_scale_nan.py"),
         str(export_dir)],
        capture_output=True, text=True,
    )
    print(scan.stdout.strip())
    if scan.returncode != 0:
        raise SystemExit("E4M3 NaN bytes found in weight scales — requantize on modelopt >= 0.45")

    # 4. Tokenizer round-trip vs source
    from transformers import AutoTokenizer

    tok_src = AutoTokenizer.from_pretrained(str(source), trust_remote_code=True)
    tok_exp = AutoTokenizer.from_pretrained(str(export_dir), trust_remote_code=True)
    if type(tok_exp).__name__ != "Llmjp4Tokenizer":
        raise SystemExit(f"wrong tokenizer class: {type(tok_exp).__name__}")
    mismatches = 0
    for s in ROUND_TRIP_STRINGS:
        ids_src = tok_src.encode(s, add_special_tokens=False)
        ids_exp = tok_exp.encode(s, add_special_tokens=False)
        dec_exp = tok_exp.decode(ids_exp)
        if ids_src != ids_exp or dec_exp != tok_src.decode(ids_src):
            mismatches += 1
            print(f"  MISMATCH: {s!r}")
    if mismatches:
        raise SystemExit(f"{mismatches} round-trip mismatches")
    print(f"tokenizer round-trip passed on {len(ROUND_TRIP_STRINGS)} strings "
          f"(class {type(tok_exp).__name__})")

    print(
        "\nPatch complete. Final load test is serving (transformers cannot run "
        "NVFP4 natively):\n"
        f"  vllm serve {export_dir} --trust-remote-code --reasoning-parser llmjp4\n"
        "then check startup logs for quantization=modelopt_mixed and the "
        "FlashInferCutlassNvFp4LinearKernel selection, and smoke-test "
        "/v1/chat/completions."
    )


if __name__ == "__main__":
    main()