#!/usr/bin/env python3
"""Export all four catalog models from checkpoints synced by sync_upstream_assets.py."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import subprocess
import sys
import tempfile


ROOT = Path(__file__).resolve().parents[1]
BASE_REVISION = "f3ff3571791e39611d31c381e3a41a3af07b4987"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--checkpoint-manifest",
        type=Path,
        default=Path(tempfile.gettempdir()) / "kokoro-kmp-checkpoints.json",
    )
    args = parser.parse_args()
    if not args.checkpoint_manifest.is_file():
        raise SystemExit(
            f"Missing {args.checkpoint_manifest}; run scripts/sync_upstream_assets.py first"
        )
    checkpoints = json.loads(args.checkpoint_manifest.read_text(encoding="utf-8"))

    exports = [
        {
            "id": "kokoro-v1.0-fp32",
            "checkpointName": "kokoro-v1_0.pth",
            "config": "runtime/kokoro-v1.0-config.json",
            "voice": "voices/en-us/af_heart.bin",
            "output": "models/standard/kokoro-v1.0-fp32.onnx",
            "language": "multilingual-v1.0",
            "phonemes": "həlˈO wˈɜɹld.",
            "repository": "hexgrad/Kokoro-82M",
            "revision": BASE_REVISION,
        },
        {
            "id": "kokoro-v1.1-zh-fp32",
            "checkpointName": "kokoro-v1_1-zh.pth",
            "config": "runtime/kokoro-v1.1-zh-config.json",
            "voice": "voices/zh/zf_001.bin",
            "output": "models/zh/kokoro-v1.1-zh-fp32.onnx",
            "language": "zh-en-v1.1",
            "phonemes": "ㄋㄧ2ㄏㄠ3, ㄕ十4ㄐㄝ4.",
            "repository": "hexgrad/Kokoro-82M-v1.1-zh",
            "revision": "01e7505bd6a7a2ac4975463114c3a7650a9f7218",
        },
        {
            "id": "kokoro-pl-fp32",
            "checkpointName": "finetunes/kokoro_polish_converted.pth",
            "config": "runtime/kokoro-v1.0-config.json",
            "voice": "voices/pl/pm_mateusz.bin",
            "output": "models/pl/kokoro-pl-fp32.onnx",
            "language": "pl",
            "phonemes": "zˈaʒuwtɕ ɡˈɛɲɕlɔ̃ jˈaʒɲ.",
            "repository": "software-mansion/react-native-executorch-kokoro",
            "revision": "3744b57964eab7df6e8c48f0b84badb29e14df07",
        },
        {
            "id": "kokoro-de-fp32",
            "checkpointName": "finetunes/kokoro_german_converted.pth",
            "config": "runtime/kokoro-v1.0-config.json",
            "voice": "voices/de/df_anna.bin",
            "output": "models/de/kokoro-de-fp32.onnx",
            "language": "de",
            "phonemes": "hˈaloː vˈɛlt.",
            "repository": "software-mansion/react-native-executorch-kokoro",
            "revision": "3744b57964eab7df6e8c48f0b84badb29e14df07",
        },
    ]

    exporter = ROOT / "scripts/export_onnx.py"
    for item in exports:
        checkpoint = checkpoints.get(item["id"])
        if checkpoint is None:
            raise RuntimeError(f"Checkpoint manifest has no {item['id']}")
        command = [
            sys.executable,
            str(exporter),
            "--checkpoint",
            checkpoint,
            "--source-checkpoint-name",
            item["checkpointName"],
            "--config",
            item["config"],
            "--voice",
            item["voice"],
            "--output",
            item["output"],
            "--language",
            item["language"],
            "--reference-phonemes",
            item["phonemes"],
            "--source-repository",
            item["repository"],
            "--source-revision",
            item["revision"],
            "--base-revision",
            BASE_REVISION,
        ]
        print(f"Exporting {item['id']}", flush=True)
        subprocess.run(command, check=True, cwd=ROOT)


if __name__ == "__main__":
    main()