#!/usr/bin/env python3
"""
Download all ONNX models for Hebrew Unified NLP

Usage:
    python download_models.py              # Download all 36 voices
    python download_models.py --essential  # Download essential only (phonikud + Hebrew + ryan-high)
"""

import sys
import shutil
from pathlib import Path
from huggingface_hub import hf_hub_download

ONNX_DIR = Path("./onnx")
PIPER_DIR = ONNX_DIR / "piper-voices"

# All available English voices
EN_VOICES = {
    # en_GB
    "en_GB-alan-low": ("en_GB", "alan", "low"),
    "en_GB-alan-medium": ("en_GB", "alan", "medium"),
    "en_GB-alba-medium": ("en_GB", "alba", "medium"),
    "en_GB-aru-medium": ("en_GB", "aru", "medium"),
    "en_GB-cori-medium": ("en_GB", "cori", "medium"),
    "en_GB-cori-high": ("en_GB", "cori", "high"),
    "en_GB-jenny_dioco-medium": ("en_GB", "jenny_dioco", "medium"),
    "en_GB-northern_english_male-medium": ("en_GB", "northern_english_male", "medium"),
    "en_GB-semaine-medium": ("en_GB", "semaine", "medium"),
    "en_GB-southern_english_female-low": ("en_GB", "southern_english_female", "low"),
    "en_GB-vctk-medium": ("en_GB", "vctk", "medium"),
    # en_US
    "en_US-amy-low": ("en_US", "amy", "low"),
    "en_US-amy-medium": ("en_US", "amy", "medium"),
    "en_US-arctic-medium": ("en_US", "arctic", "medium"),
    "en_US-bryce-medium": ("en_US", "bryce", "medium"),
    "en_US-danny-low": ("en_US", "danny", "low"),
    "en_US-hfc_female-medium": ("en_US", "hfc_female", "medium"),
    "en_US-hfc_male-medium": ("en_US", "hfc_male", "medium"),
    "en_US-joe-medium": ("en_US", "joe", "medium"),
    "en_US-john-medium": ("en_US", "john", "medium"),
    "en_US-kathleen-low": ("en_US", "kathleen", "low"),
    "en_US-kristin-medium": ("en_US", "kristin", "medium"),
    "en_US-kusal-medium": ("en_US", "kusal", "medium"),
    "en_US-l2arctic-medium": ("en_US", "l2arctic", "medium"),
    "en_US-lessac-high": ("en_US", "lessac", "high"),
    "en_US-lessac-low": ("en_US", "lessac", "low"),
    "en_US-lessac-medium": ("en_US", "lessac", "medium"),
    "en_US-libritts-high": ("en_US", "libritts", "high"),
    "en_US-libritts_r-medium": ("en_US", "libritts_r", "medium"),
    "en_US-ljspeech-high": ("en_US", "ljspeech", "high"),
    "en_US-ljspeech-medium": ("en_US", "ljspeech", "medium"),
    "en_US-norman-medium": ("en_US", "norman", "medium"),
    "en_US-ryan-high": ("en_US", "ryan", "high"),
    "en_US-ryan-low": ("en_US", "ryan", "low"),
    "en_US-ryan-medium": ("en_US", "ryan", "medium"),
}


def download_file(repo_id: str, filename: str, dest_path: Path):
    """Download a file from HuggingFace Hub to destination."""
    if dest_path.exists():
        print(f"  [skip] {dest_path.name} already exists")
        return False

    print(f"  [download] {filename}...")
    downloaded = hf_hub_download(repo_id=repo_id, filename=filename)
    dest_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy(downloaded, dest_path)
    print(f"  [done] {dest_path.name}")
    return True


def download_phonikud():
    """Download Phonikud ONNX model."""
    print("\n=== Phonikud Model ===")
    download_file(
        "thewh1teagle/phonikud-onnx",
        "phonikud-1.0.int8.onnx",
        ONNX_DIR / "phonikud-1.0.int8.onnx"
    )


def download_hebrew_voice():
    """Download Hebrew Piper voice."""
    print("\n=== Hebrew Voice (he_IL-phonikud) ===")
    download_file(
        "thewh1teagle/phonikud-tts-checkpoints",
        "model.onnx",
        PIPER_DIR / "he_IL-phonikud.onnx"
    )
    download_file(
        "thewh1teagle/phonikud-tts-checkpoints",
        "model.config.json",
        PIPER_DIR / "he_IL-phonikud.onnx.json"
    )


def download_english_voice(voice_name: str):
    """Download an English Piper voice."""
    region, speaker, quality = EN_VOICES[voice_name]
    subdir = f"en/{region}/{speaker}/{quality}"

    download_file(
        "rhasspy/piper-voices",
        f"{subdir}/{voice_name}.onnx",
        PIPER_DIR / f"{voice_name}.onnx"
    )
    download_file(
        "rhasspy/piper-voices",
        f"{subdir}/{voice_name}.onnx.json",
        PIPER_DIR / f"{voice_name}.onnx.json"
    )


def main():
    essential_only = "--essential" in sys.argv

    print("=" * 50)
    print("Hebrew Unified NLP - Model Downloader")
    print("=" * 50)

    # Create directories
    PIPER_DIR.mkdir(parents=True, exist_ok=True)

    # Essential downloads
    download_phonikud()
    download_hebrew_voice()

    print("\n=== English Voice (en_US-ryan-high) ===")
    download_english_voice("en_US-ryan-high")

    if not essential_only:
        print("\n=== All English Voices ===")
        for voice in EN_VOICES:
            if voice != "en_US-ryan-high":  # Already downloaded
                print(f"\n  [{voice}]")
                download_english_voice(voice)

    # Summary
    print("\n" + "=" * 50)
    print("Download Complete!")
    print("=" * 50)

    total_size = sum(f.stat().st_size for f in ONNX_DIR.rglob("*") if f.is_file())
    print(f"\nLocation: {ONNX_DIR.absolute()}")
    print(f"Total size: {total_size / 1024 / 1024:.1f} MB")

    files = list(PIPER_DIR.glob("*.onnx"))
    print(f"Voices: {len(files)} ({len(files) - 1} English + 1 Hebrew)")


if __name__ == "__main__":
    main()