#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.9"
# dependencies = ["huggingface_hub"]
# ///
"""Download crane TTS, ASR, VAD, and LLM models from Hugging Face.

Examples:
    ./crane-model-download --list
    ./crane-model-download --model voxtral --path /srv/models
    ./crane-model-download --model qwen3-tts-customvoice-1.7b --path ~/models --token hf_...
    ./crane-model-download --model qwen3-asr-0.6b --path ~/models
    ./crane-model-download --model silero-vad --path ~/models
    ./crane-model-download --model qwen3.5-9b-gguf --path ~/models --quant Q8_0
    ./crane-model-download --model g2p-lexicons --path ~/models
    ./crane-model-download --model kokoro --path ~/models

Run directly (uv picks up the inline metadata above and manages the
`huggingface_hub` dependency in an ephemeral venv, no setup needed):

    uv run tools/crane-model-download --list

Without `uv`, install the dependency yourself and run with plain python3:

    pip install huggingface_hub
    python3 tools/crane-model-download --list

huggingface_hub >= 0.32.0 bundles the `hf-xet` accelerator by default,
which speeds up downloads automatically.
"""

import argparse
import os
import shutil
import sys
import urllib.error
import urllib.request

MODELS = {
    "qwen3-tts-base": {
        "repo_id": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
        "dirname": "Qwen3-TTS-12Hz-0.6B-Base",
        "kind": "tts",
        "description": "Qwen3-TTS 0.6B, voice cloning via reference audio (Apache-2.0)",
    },
    "qwen3-tts-customvoice-0.6b": {
        "repo_id": "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
        "dirname": "Qwen3-TTS-12Hz-0.6B-CustomVoice",
        "kind": "tts",
        "description": "Qwen3-TTS 0.6B, predefined speakers (Apache-2.0)",
    },
    "qwen3-tts-customvoice-1.7b": {
        "repo_id": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
        "dirname": "Qwen3-TTS-12Hz-1.7B-CustomVoice",
        "kind": "tts",
        "description": "Qwen3-TTS 1.7B, predefined speakers (Apache-2.0)",
    },
    "voxtral": {
        "repo_id": "mistralai/Voxtral-4B-TTS-2603",
        "dirname": "Voxtral-4B-TTS-2603",
        "kind": "tts",
        "description": "Voxtral 4B, 10 languages / 20 voices (CC-BY-NC-4.0)",
    },
    "qwen3-asr-0.6b": {
        "repo_id": "Qwen/Qwen3-ASR-0.6B-hf",
        "dirname": "Qwen3-ASR-0.6B-hf",
        "kind": "asr",
        "description": "Qwen3-ASR 0.6B, speech-to-text (Apache-2.0)",
    },
    "qwen3-asr-1.7b": {
        "repo_id": "Qwen/Qwen3-ASR-1.7B-hf",
        "dirname": "Qwen3-ASR-1.7B-hf",
        "kind": "asr",
        "description": "Qwen3-ASR 1.7B, speech-to-text (Apache-2.0)",
    },
    "silero-vad": {
        "repo_id": "onnx-community/silero-vad",
        "dirname": "silero-vad",
        "kind": "vad",
        "files": {"onnx/model.onnx": "model.onnx"},
        "description": "Silero VAD, voice activity detection (MIT)",
    },
    "qwen3.5-0.8b-gguf": {
        "repo_id": "unsloth/Qwen3.5-0.8B-GGUF",
        "dirname": "Qwen3.5-0.8B-GGUF",
        "kind": "llm",
        "quant_template": "Qwen3.5-0.8B-{quant}.gguf",
        "description": "Qwen 3.5 0.8B, hybrid GDN + attention, GGUF quant (Apache-2.0)",
    },
    "qwen3.5-4b-gguf": {
        "repo_id": "unsloth/Qwen3.5-4B-GGUF",
        "dirname": "Qwen3.5-4B-GGUF",
        "kind": "llm",
        "quant_template": "Qwen3.5-4B-{quant}.gguf",
        "description": "Qwen 3.5 4B, hybrid GDN + attention, GGUF quant (Apache-2.0)",
    },
    "qwen3.5-9b-gguf": {
        "repo_id": "unsloth/Qwen3.5-9B-GGUF",
        "dirname": "Qwen3.5-9B-GGUF",
        "kind": "llm",
        "quant_template": "Qwen3.5-9B-{quant}.gguf",
        "description": "Qwen 3.5 9B, hybrid GDN + attention, GGUF quant (Apache-2.0)",
    },
    "qwen3.8-27b-gguf": {
        "repo_id": "unsloth/Qwen3.8-27B-GGUF",
        "dirname": "Qwen3.8-27B-GGUF",
        "kind": "llm",
        "quant_template": "Qwen3.8-27B-{quant}.gguf",
        "description": "Qwen 3.8 27B, same arch as Qwen 3.5 scaled up, GGUF quant (Apache-2.0)",
    },
    "g2p-lexicons": {
        "repo_id": "crane-local-ai/g2p-lexicons",
        "dirname": "g2p-lexicons",
        "kind": "g2p",
        "repo_type": "dataset",
        # Local names match what MoonshineG2p::from_g2p_dir hardcodes
        # (en_us/dict_filtered_heteronyms.tsv, de/dict.tsv), replacing the
        # lexicons moonshine-g2p-en-us/de used to fetch from Moonshine's CDN.
        "files": {
            "de/de.tsv": "de/dict.tsv",
            "de/PROVENANCE.md": "de/PROVENANCE.md",
            "en_US/en_US.tsv": "en_us/dict_filtered_heteronyms.tsv",
            "en_US/PROVENANCE.md": "en_us/PROVENANCE.md",
        },
        "description": "Crane G2P word-to-IPA lexicons: German (Wiktionary, CC BY-SA 4.0) + English/US (ipa-dict, MIT)",
    },
    "g2p-oov-en-us": {
        "repo_id": "moonshine-ai/moonshine",
        "dirname": "g2p-lexicons",
        "kind": "g2p",
        "source": "url",
        "base_url": "https://download.moonshine.ai/tts/en_us",
        # Lands in the same g2p-lexicons/en_us/ directory as the g2p-lexicons
        # entry above -- our own dataset has no OOV model, so English still
        # needs this neural fallback from Moonshine.
        "files": {
            "g2p-config.json": "en_us/g2p-config.json",
            "oov/model.onnx": "en_us/oov/model.onnx",
            "oov/onnx-config.json": "en_us/oov/onnx-config.json",
        },
        "description": "Moonshine English G2P OOV fallback model only, no lexicon -- use with g2p-lexicons (MIT)",
    },
    "kokoro": {
        "repo_id": "onnx-community/Kokoro-82M-v1.0-ONNX",
        "dirname": "Kokoro-82M-v1.0-ONNX",
        "kind": "tts",
        "description": "Kokoro 82M, ONNX, 55 voices / 10 languages (Apache-2.0)",
    },
    "kokoro-de": {
        "repo_id": "crane-local-ai/Kokoro-82M-v1.0-German-ONNX",
        "dirname": "Kokoro-82M-v1.0-German-ONNX",
        "kind": "tts",
        "description": "Kokoro 82M, ONNX, German-only fine-tune, 1 voice (Apache-2.0)",
    },
    "minicpm5-1b": {
        "repo_id": "openbmb/MiniCPM5-1B",
        "dirname": "MiniCPM5-1B",
        "kind": "llm",
        "description": "MiniCPM5 1B, plain dense chat model (Apache-2.0)",
    },
    "minicpm5-1b-gguf": {
        "repo_id": "openbmb/MiniCPM5-1B-GGUF",
        "dirname": "MiniCPM5-1B-GGUF",
        "kind": "llm",
        "quant_template": "MiniCPM5-1B-{quant}.gguf",
        "description": "MiniCPM5 1B, GGUF quant with embedded tokenizer (Apache-2.0)",
    },
    "minicpm-v-4.6": {
        "repo_id": "openbmb/MiniCPM-V-4.6",
        "dirname": "MiniCPM-V-4.6",
        "kind": "vlm",
        "description": "MiniCPM-V 4.6, vision-language chat (Apache-2.0)",
    },
    "minicpm-o-4.5": {
        "repo_id": "openbmb/MiniCPM-o-4_5",
        "dirname": "MiniCPM-o-4_5",
        "kind": "omni",
        "description": "MiniCPM-o 4.5, full omni: vision + audio understanding, speech-token TTS, full-duplex live audio chat (Apache-2.0)",
    },
    "minicpm-o-4.5-llm-gguf": {
        "repo_id": "openbmb/MiniCPM-o-4_5-gguf",
        "dirname": "MiniCPM-o-4_5-gguf",
        "kind": "omni",
        "quant_template": "MiniCPM-o-4_5-{quant}.gguf",
        "description": "MiniCPM-o 4.5 standalone LLM-tower GGUF, for --llm-gguf (genuinely quantized, cuts VRAM ~19.8GB to ~11.2GB at Q8_0) (Apache-2.0)",
    },
    "minicpm-o-4.5-other-ggufs": {
        "repo_id": "openbmb/MiniCPM-o-4_5-gguf",
        "dirname": "MiniCPM-o-4_5-gguf",
        "kind": "omni",
        "files": {
            "audio/MiniCPM-o-4_5-audio-F16.gguf": "MiniCPM-o-4_5-audio-F16.gguf",
            "vision/MiniCPM-o-4_5-vision-F16.gguf": "MiniCPM-o-4_5-vision-F16.gguf",
            "tts/MiniCPM-o-4_5-tts-F16.gguf": "MiniCPM-o-4_5-tts-F16.gguf",
        },
        "description": "MiniCPM-o 4.5 audio/vision/tts tower GGUFs (F16 only, no VRAM savings -- pure-GGUF deployment convenience) (Apache-2.0)",
    },
    "voxcpm2": {
        "repo_id": "openbmb/VoxCPM2",
        "dirname": "VoxCPM2",
        "kind": "tts",
        # `openbmb/VoxCPM2` only ships `audiovae.pth` (a PyTorch pickle);
        # candle has no pickle reader, so Crane needs `audiovae.safetensors`
        # instead. Rather than making every user run the one-time conversion
        # themselves, `extra_files` pulls an already-converted copy (same
        # weights, straight `torch.load` -> `safetensors.torch.save_file`,
        # nothing else changed) from a second repo into the same directory.
        "extra_files": {
            "hahihula/VoxCPM2-audiovae-safetensors": {
                "audiovae.safetensors": "audiovae.safetensors",
            },
        },
        "description": "VoxCPM2, tokenizer-free diffusion TTS, voice cloning + prompt cache (Apache-2.0)",
    },
    "muscriptor-small": {
        "repo_id": "MuScriptor/muscriptor-small",
        "dirname": "muscriptor-small",
        "kind": "amt",
        "description": "MuScriptor small (~100M), audio-to-MIDI transcription, fastest/lightest variant (CC-BY-NC-4.0, gated)",
    },
    "muscriptor-medium": {
        "repo_id": "MuScriptor/muscriptor-medium",
        "dirname": "muscriptor-medium",
        "kind": "amt",
        "description": "MuScriptor medium (~300M), audio-to-MIDI transcription, quality/speed trade-off (CC-BY-NC-4.0, gated)",
    },
    "muscriptor-large": {
        "repo_id": "MuScriptor/muscriptor-large",
        "dirname": "muscriptor-large",
        "kind": "amt",
        "description": "MuScriptor large (~1.3B), audio-to-MIDI transcription, best quality (CC-BY-NC-4.0, gated)",
    },
}

_VALID_KINDS = ("tts", "asr", "vad", "llm", "g2p", "vlm", "omni", "amt")
_REQUIRED_KEYS = ("repo_id", "dirname", "kind", "description")
for _key, _info in MODELS.items():
    _missing = [k for k in _REQUIRED_KEYS if k not in _info]
    if _missing:
        raise ValueError(f"MODELS[{_key!r}] is missing required key(s): {_missing}")
    if _info["kind"] not in _VALID_KINDS:
        raise ValueError(
            f"MODELS[{_key!r}] has invalid kind {_info['kind']!r}; "
            f"expected one of {_VALID_KINDS}"
        )


def list_models():
    print("Available models:\n")
    width = max(len(key) for key in MODELS)
    for key, info in MODELS.items():
        kind = info["kind"].upper()
        print(f"  {key:<{width}}  [{kind}]  {info['description']}")


def download_url_files(info, local_dir):
    base = info["base_url"]
    for rel_path, local_name in info["files"].items():
        url = f"{base}/{rel_path}"
        target = os.path.join(local_dir, local_name)
        os.makedirs(os.path.dirname(target), exist_ok=True)
        print(f"  {url}")
        # download.moonshine.ai's CDN returns 403 for the default
        # Python-urllib/x.y User-Agent; send a browser-like one instead.
        request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        try:
            with urllib.request.urlopen(request) as response, open(target, "wb") as f:
                shutil.copyfileobj(response, f)
        except urllib.error.HTTPError as e:
            print(f"error: download failed for {url}: {e}", file=sys.stderr)
            sys.exit(1)


def print_done(info, target):
    if info["kind"] == "amt":
        # MuScriptor isn't served through crane-serve — it's a standalone
        # CLI, so the generic crane-serve hint below would be actively
        # wrong here rather than just imprecise.
        print(
            f"\nDone. Use with: cargo run -p crane-examples --release --features cuda "
            f"--bin muscriptor_transcribe -- --model-dir {target} --transcribe audio.wav"
        )
    elif info["kind"] == "g2p":
        # g2p assets aren't a --model-path themselves; Kokoro expects them
        # symlinked in at {model_path}/g2p (see
        # crane_core::models::g2p::MoonshineG2p::from_g2p_dir).
        print(f"\nDone. Use with: ln -s {target} <model_path>/g2p")
    else:
        print(f"\nDone. Use with: crane-serve --model-path {target}")


def download_model(model_key, path, token, quant):
    info = MODELS[model_key]
    local_dir = os.path.join(os.path.expanduser(path), info["kind"], info["dirname"])

    if info.get("source") == "url":
        print(f"Downloading {info['repo_id']} ({info['base_url']}) to {local_dir} ...")
        os.makedirs(local_dir, exist_ok=True)
        download_url_files(info, local_dir)
        print_done(info, local_dir)
        return

    try:
        from huggingface_hub import hf_hub_download, snapshot_download
        from huggingface_hub.utils import HfHubHTTPError
    except ImportError:
        print(
            "error: huggingface_hub is not installed. Install it with:\n\n"
            "    pip install huggingface_hub\n",
            file=sys.stderr,
        )
        sys.exit(1)

    quant_template = info.get("quant_template")
    repo_type = info.get("repo_type", "model")

    if quant_template:
        filename = quant_template.format(quant=quant)
        target = os.path.join(local_dir, filename)
        print(f"Downloading {info['repo_id']}:{filename} to {target} ...")
    else:
        target = local_dir
        print(f"Downloading {info['repo_id']} to {local_dir} ...")

    try:
        if quant_template:
            os.makedirs(local_dir, exist_ok=True)
            cached = hf_hub_download(
                repo_id=info["repo_id"], filename=filename, token=token
            )
            shutil.copy2(cached, target)
        else:
            files = info.get("files")
            if files:
                os.makedirs(local_dir, exist_ok=True)
                for repo_path, local_name in files.items():
                    cached = hf_hub_download(
                        repo_id=info["repo_id"],
                        filename=repo_path,
                        repo_type=repo_type,
                        token=token,
                    )
                    local_target = os.path.join(local_dir, local_name)
                    os.makedirs(os.path.dirname(local_target), exist_ok=True)
                    shutil.copy2(cached, local_target)
            else:
                snapshot_download(
                    repo_id=info["repo_id"],
                    local_dir=local_dir,
                    repo_type=repo_type,
                    token=token,
                )

        # Additional files pulled from a *different* repo into the same
        # local_dir (e.g. a pre-converted weight file the primary repo
        # doesn't ship) — downloaded after the primary content above so a
        # failure here doesn't leave the main download half-finished.
        extra_files = info.get("extra_files")
        if extra_files:
            os.makedirs(local_dir, exist_ok=True)
            for extra_repo_id, repo_files in extra_files.items():
                for repo_path, local_name in repo_files.items():
                    print(f"Downloading {extra_repo_id}:{repo_path} ...")
                    cached = hf_hub_download(
                        repo_id=extra_repo_id, filename=repo_path, token=token
                    )
                    shutil.copy2(cached, os.path.join(local_dir, local_name))
    except HfHubHTTPError as e:
        status = e.response.status_code if e.response is not None else None
        print(f"error: download failed: {e}", file=sys.stderr)
        if status in (401, 403):
            print(
                "\nThis model may be gated. Make sure you've accepted its "
                "license on the Hugging Face model page, then pass a valid "
                "--token (or set HF_TOKEN).",
                file=sys.stderr,
            )
        elif status == 404 and quant_template:
            print(
                f"\nNo file named {filename!r} in {info['repo_id']}. Check "
                "the repo's file list on Hugging Face for valid --quant "
                "suffixes (e.g. Q4_K_M, Q5_K_M, Q8_0, UD-Q4_K_XL).",
                file=sys.stderr,
            )
        sys.exit(1)

    print_done(info, target)


def main():
    parser = argparse.ArgumentParser(
        description="Download TTS, ASR, and VAD models for crane-serve from Hugging Face.",
        epilog="examples:\n"
        "  crane-model-download --list\n"
        "  crane-model-download --model voxtral --path /srv/models\n"
        "  crane-model-download --model qwen3-tts-customvoice-1.7b --path ~/models --token hf_...\n"
        "  crane-model-download --model qwen3-asr-0.6b --path ~/models\n"
        "  crane-model-download --model silero-vad --path ~/models\n"
        "  crane-model-download --model qwen3.5-9b-gguf --path ~/models --quant Q8_0\n"
        "  crane-model-download --model g2p-lexicons --path ~/models\n"
        "  crane-model-download --model g2p-oov-en-us --path ~/models\n"
        "  crane-model-download --model kokoro --path ~/models\n",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--list", action="store_true", help="List available models and exit"
    )
    group.add_argument(
        "--model", choices=list(MODELS.keys()), help="Model to download"
    )
    parser.add_argument(
        "--path", help="Parent directory to download the model into"
    )
    parser.add_argument(
        "--token",
        help="Hugging Face auth token (prefer the HF_TOKEN env var instead, "
        "to avoid leaking it via shell history or `ps`)",
    )
    parser.add_argument(
        "--quant",
        default="Q4_K_M",
        help="Quantization suffix to download, for GGUF models only "
        "(e.g. Q4_K_M, Q5_K_M, Q8_0, UD-Q4_K_XL). Default: Q4_K_M",
    )
    args = parser.parse_args()

    if args.list:
        list_models()
        return

    if not args.path:
        parser.error("--path is required when using --model")

    download_model(args.model, args.path, args.token, args.quant)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)
