"""Prepare and audit a separate Swift baseline; never mutate the source checkpoint."""
import argparse
import copy
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

from common import ROOT, RUN, SOURCE, BASELINE, read_json, write_json, sha256, stamp


def clone(source, destination):
    source, destination = Path(source), Path(destination)
    destination.mkdir(parents=True, exist_ok=True)
    for p in source.iterdir():
        if not p.is_file() or ".bak" in p.name or p.name.endswith(".tmp"):
            continue
        target = destination / p.name
        if target.exists():
            continue
        if p.suffix == ".safetensors":
            os.link(p, target)
        else:
            shutil.copy2(p, target)


def audit(directory):
    from safetensors import safe_open
    directory = Path(directory)
    index = read_json(directory / "model.safetensors.index.json")
    mapping = index["weight_map"]
    actual = {}
    duplicates = []
    total = 0
    for name in sorted(set(mapping.values())):
        p = directory / name
        total += p.stat().st_size
        with safe_open(p, framework="pt") as f:
            for key in f.keys():
                if key in actual:
                    duplicates.append(key)
                actual[key] = name
    missing = sorted(set(mapping) - set(actual))
    orphaned = sorted(set(actual) - set(mapping))
    wrong = [k for k in mapping if actual.get(k) != mapping[k]]
    assert not (duplicates or missing or orphaned or wrong), (duplicates, missing, orphaned, wrong)
    config = read_json(directory / "config.json")
    result = {"time": stamp(), "directory": str(directory), "tensor_count": len(actual),
              "shard_bytes": total, "duplicate_keys": duplicates,
              "quantization": config["quantization_config"]}
    write_json(directory / "hyperqwen-audit.json", result)
    print("Verified", len(actual), "unique indexed tensors;", round(total / 2**30, 2), "GiB", flush=True)
    return result


def baseline():
    if (BASELINE / "hyperqwen-build.json").exists():
        return audit(BASELINE)
    clone(SOURCE, BASELINE)
    index = read_json(BASELINE / "model.safetensors.index.json")
    if "lm_head.weight" in index["weight_map"]:
        subprocess.run([sys.executable, str(ROOT / "prepare/quant_heads_stream.py"), str(BASELINE)], check=True)
    if not (BASELINE / "mtp_draft_vocab_ids.pt").exists():
        subprocess.run([sys.executable, str(ROOT / "prepare/build_draft_vocab.py"), str(BASELINE),
                        "--ids", str(ROOT / "prepare/draft_vocab_ids.json")], check=True)
    # config.json is authoritative. Some exports carry a second quantization file.
    config = read_json(BASELINE / "config.json")
    if (BASELINE / "quantization_config.json").exists():
        write_json(BASELINE / "quantization_config.json", config["quantization_config"])
    write_json(BASELINE / "hyperqwen-build.json", {
        "created": stamp(), "source": read_json(RUN / "source.json"),
        "variant": "baseline", "body": "unchanged asymmetric AWQ INT4 group128",
        "embedding_bits": 8, "lm_head_bits": 8, "mtp_bits": 8,
        "preparation_script_sha256": sha256(ROOT / "prepare/quant_heads_stream.py"),
        "draft_vocabulary": "HyperQwen reference ids; rebuilt for Swift in fast variant",
    })
    audit(BASELINE)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("action", choices=["baseline", "audit"])
    ap.add_argument("--model", type=Path, default=BASELINE)
    args = ap.parse_args()
    baseline() if args.action == "baseline" else audit(args.model)