"""Fix tokenizer v3: add 6 ChatML tokens (32000-32005) to model.vocab."""
import json
import shutil
from pathlib import Path

src = Path("tokenizer.json")
backup = Path("tokenizer.json.bak")
if not backup.exists():
    shutil.copy(src, backup)
    print(f"Backup: {backup}")

with open(src) as f:
    t = json.load(f)

vocab = t["model"]["vocab"]
print(f"vocab before: {len(vocab)}")

# Add 6 ChatML tokens at their explicit IDs
chatml_specials = {
    32000: "<|im_start|>",
    32001: "<|im_end|>",
    32002: "<|assistant|>",
    32003: "<|tool_call|>",
    32004: "<|tool_result|>",
    32005: "<|task_type|>",
}
added = 0
for vid, content in chatml_specials.items():
    if content not in vocab:
        vocab[content] = vid
        added += 1
print(f"vocab after: {len(vocab)} (+{added})")

# Re-check coherence
inv = {v: k for k, v in vocab.items()}
mismatches = []
for tok in t.get("added_tokens", []):
    tid = tok["id"]
    if inv.get(tid) != tok["content"]:
        mismatches.append((tid, tok["content"], inv.get(tid)))
print(f"Remaining mismatches: {len(mismatches)}")
for m in mismatches:
    print(f"  {m}")

with open(src, "w", encoding="utf-8") as f:
    json.dump(t, f, ensure_ascii=False)
print(f"Saved {src}")

# Test load
from tokenizers import Tokenizer
try:
    tk = Tokenizer.from_file(str(src))
    print(f"LOAD OK vocab_size={tk.get_vocab_size()}")
    # Test encode/decode
    s = "Hello [trading-specialist], <|im_start|>user content<|im_end|>"
    ids = tk.encode(s).ids
    dec = tk.decode(ids)
    print(f"encode test ids[:10]={ids[:10]}")
    print(f"decode test: {dec!r}")
except Exception as e:
    print(f"LOAD FAIL: {e}")