#!/usr/bin/env python3
"""Builds ASRModels.bundle for SpeechKit distribution.

Collects models + assets from the iphone-asr workspace, precompiles
mlpackage -> mlmodelc (xcrun coremlcompiler), computes SHA-256 for every
asset, and writes a schema-v2 manifest.json.

Usage: python3 make_asset_bundle.py [output_dir]
"""
import hashlib
import json
import shutil
import subprocess
import sys
from pathlib import Path

WORKSPACE = Path(__file__).resolve().parents[2]  # iphone-asr/
OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else WORKSPACE / "dist/ASRModels.bundle"
BUNDLE_VERSION = "1.0.0"

SOURCES = {
    # logical name -> (source path, dest file name)
    "tower": (WORKSPACE / "artifacts/audio_tower_multifunction_w8.mlpackage", "tower.mlpackage"),
    "lm_prefill": (WORKSPACE / "artifacts/lm_shared/lm_cache_prefill_int4.onnx", "lm_prefill.onnx"),
    "lm_decode": (WORKSPACE / "artifacts/lm_shared/lm_cache_decode_int4.onnx", "lm_decode.onnx"),
    "lm_shared_data": (WORKSPACE / "artifacts/lm_shared/lm_shared_int4.data", "lm_shared_int4.data"),
    "mask_gen": (WORKSPACE / "artifacts/mask_gen.onnx", "mask_gen.onnx"),
    "asr_manifest": (WORKSPACE / "ios_assets/manifest.json", "asr_manifest.json"),
    "hann_window": (WORKSPACE / "ios_assets/hann_window.bin", "hann_window.bin"),
    "mel_filters": (WORKSPACE / "ios_assets/mel_filters.bin", "mel_filters.bin"),
    "projector_norm_w": (WORKSPACE / "ios_assets/projector_norm_w.bin", "projector_norm_w.bin"),
    "projector_norm_b": (WORKSPACE / "ios_assets/projector_norm_b.bin", "projector_norm_b.bin"),
    "projector_lin_w": (WORKSPACE / "ios_assets/projector_lin_w.bin", "projector_lin_w.bin"),
    "projector_lin_b": (WORKSPACE / "ios_assets/projector_lin_b.bin", "projector_lin_b.bin"),
    "token_embedding": (WORKSPACE / "ios_assets/token_embedding_fp16.bin", "token_embedding_fp16.bin"),
    "vocab_bytes": (WORKSPACE / "ios_assets/vocab_bytes.bin", "vocab_bytes.bin"),
    "vocab_offsets": (WORKSPACE / "ios_assets/vocab_offsets.bin", "vocab_offsets.bin"),
}


def sha256_path(p: Path) -> str:
    h = hashlib.sha256()
    if p.is_dir():
        for f in sorted(p.rglob("*")):
            if f.is_file():
                h.update(f.relative_to(p).as_posix().encode())
                h.update(f.read_bytes())
    else:
        with p.open("rb") as fh:
            for chunk in iter(lambda: fh.read(1 << 23), b""):
                h.update(chunk)
    return h.hexdigest()


def dir_size(p: Path) -> int:
    if p.is_file():
        return p.stat().st_size
    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file())


if OUT.exists():
    shutil.rmtree(OUT)
OUT.mkdir(parents=True)

assets = {}
for logical, (src, dest_name) in SOURCES.items():
    assert src.exists(), f"missing source: {src}"
    dest = OUT / dest_name
    if src.suffix == ".mlpackage":
        # precompile so devices never pay the mlpackage->mlmodelc conversion
        dest = OUT / (Path(dest_name).stem + ".mlmodelc")
        subprocess.run(["xcrun", "coremlcompiler", "compile", str(src), str(OUT)],
                       check=True, capture_output=True)
        compiled = OUT / (src.stem + ".mlmodelc")
        if compiled != dest:
            compiled.rename(dest)
    elif src.is_dir():
        shutil.copytree(src, dest)
    else:
        shutil.copy2(src, dest)
    assets[logical] = {
        "file": dest.name,
        "sha256": sha256_path(dest),
        "sizeBytes": dir_size(dest),
    }
    print(f"  {logical:18s} -> {dest.name} ({dir_size(dest)/1e6:.1f}MB)")

manifest = {
    "schemaVersion": 2,
    "bundleVersion": BUNDLE_VERSION,
    "assets": assets,
    "extra": {"model": "minillm-asr-zh", "builtBy": "make_asset_bundle.py"},
}
(OUT / "manifest.json").write_text(json.dumps(manifest, indent=2))
total = sum(a["sizeBytes"] for a in assets.values())
print(f"\nbundle: {OUT}  ({total/1e6:.0f}MB, schema v2, version {BUNDLE_VERSION})")