"""Repair the reference's embedding/index mismatch in a separate evaluation copy."""
import shutil
import os
import subprocess
import sys
from common import ROOT, RUN, read_json, write_json, stamp
from checkpoint import clone, audit

SOURCE = ROOT / "models/Qwen3.8-27B-W4A16-AutoRound-fast"
TARGET = ROOT / "models/Qwen3.8-27B-W4A16-AutoRound-fast-eval"


def main():
    if (TARGET / "embedding-repair.json").exists():
        audit(TARGET)
        return
    clone(SOURCE, TARGET)
    index = read_json(TARGET / "model.safetensors.index.json")
    prefix = "model.language_model.embed_tokens"
    shard = index["weight_map"][prefix + ".weight_packed"]
    from safetensors import safe_open
    with safe_open(TARGET / shard, "pt") as file:
        assert prefix + ".weight" in file.keys()
        assert prefix + ".weight_packed" not in file.keys()
    # The legacy converter writes in place; detach this shard from source inodes.
    temporary = TARGET / (shard + ".private")
    shutil.copyfile(TARGET / shard, temporary)
    os.replace(temporary, TARGET / shard)
    for suffix in ["weight_packed", "weight_scale", "weight_shape"]:
        del index["weight_map"][prefix + "." + suffix]
    index["weight_map"][prefix + ".weight"] = shard
    write_json(TARGET / "model.safetensors.index.json", index)
    subprocess.run([sys.executable, str(ROOT / "prepare/quant_embed.py"), str(TARGET)], check=True)
    audit(TARGET)
    write_json(TARGET / "embedding-repair.json", {"created": stamp(), "source": str(SOURCE),
        "reason": "Source index declares INT8 embeddings but its shard contains BF16 embeddings.",
        "change": "Apply the repository's standard INT8 group128 embedding conversion in an isolated copy; source preserved."})


if __name__ == "__main__":
    main()