#!/usr/bin/env python3
"""Prepare a Piper ONNX export for sherpa-onnx mobile inference."""

import argparse
import json
import shutil
from pathlib import Path

import onnx


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()

    config = json.loads(args.config.read_text(encoding="utf-8"))
    output_dir: Path = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    output_model = output_dir / "aletheia_ru.onnx"
    output_config = output_dir / "aletheia_ru.onnx.json"
    output_tokens = output_dir / "tokens.txt"

    shutil.copy2(args.model, output_model)
    shutil.copy2(args.config, output_config)

    id_map = config["phoneme_id_map"]
    # sherpa-onnx's Piper lexicon maps one Unicode code point to one model ID.
    # Newer Piper configs may also contain English diphthong aliases such as
    # "aɪ". eSpeak emits their component code points, and sherpa rejects the
    # multi-code-point aliases, so only the single-code-point table is mobile-safe.
    token_rows = sorted(
        ((ids[0], symbol) for symbol, ids in id_map.items() if len(symbol) == 1),
        key=lambda row: row[0],
    )
    actual_ids = [row[0] for row in token_rows]
    if len(actual_ids) != len(set(actual_ids)):
        raise RuntimeError("phoneme_id_map contains duplicate token identifiers")
    if not actual_ids or min(actual_ids) < 0 or max(actual_ids) >= config["num_symbols"]:
        raise RuntimeError("phoneme_id_map contains a token outside the model symbol range")
    output_tokens.write_text(
        "".join(f"{symbol} {token_id}\n" for token_id, symbol in token_rows),
        encoding="utf-8",
        newline="\n",
    )

    model = onnx.load(str(output_model))
    metadata = {
        "model_type": "vits",
        "comment": "piper",
        "language": "Russian",
        "voice": config["espeak"]["voice"],
        "has_espeak": "1",
        "n_speakers": str(config["num_speakers"]),
        "sample_rate": str(config["audio"]["sample_rate"]),
    }
    existing = {item.key: item for item in model.metadata_props}
    for key, value in metadata.items():
        if key in existing:
            existing[key].value = value
        else:
            item = model.metadata_props.add()
            item.key = key
            item.value = value
    onnx.save(model, str(output_model))

    print(json.dumps({
        "model": str(output_model),
        "model_bytes": output_model.stat().st_size,
        "tokens": len(token_rows),
        "metadata": metadata,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()