#!/usr/bin/env python3
"""
Model Downloader for TypeWhisper-Nhi.
Downloads off-line speech recognition, punctuation restoration, and speaker diarization models
from Hugging Face into the models/ folder.
"""
import os
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).parent.absolute()

def download_model(repo_id, local_dir, revision=None, allowed_patterns=None):
    local_path = PROJECT_ROOT / "models" / local_dir
    print(f"\n📥 Downloading model {repo_id} -> {local_path} (revision={revision})...")
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(
            repo_id=repo_id,
            local_dir=str(local_path),
            revision=revision,
            allow_patterns=allowed_patterns
        )
        print(f"✅ Downloaded {repo_id} successfully!")
        return True
    except ImportError:
        print("❌ Error: huggingface_hub is not installed.")
        print("   Please install it first: pip install huggingface_hub")
        return False
    except Exception as e:
        print(f"❌ Error downloading {repo_id}: {e}")
        return False

def main():
    print("=" * 60)
    print("TypeWhisper-Nhi Offline Model Downloader")
    print("=" * 60)
    
    # Ensure models directory exists
    os.makedirs(PROJECT_ROOT / "models", exist_ok=True)
    
    # 1. Zipformer 30M (Lightweight RNN-T, ~120MB)
    if download_model(
        repo_id="hynt/Zipformer-30M-RNNT-6000h",
        local_dir="zipformer-30m-rnnt-6000h",
        revision="24ed30248e1c96bb690c81c24ab4e056f8cd9fce",
        allowed_patterns=["*.onnx", "tokens.txt", "bpe.model"]
    ):
        bpe_path = PROJECT_ROOT / "models" / "zipformer-30m-rnnt-6000h" / "bpe.model"
        tokens_path = PROJECT_ROOT / "models" / "zipformer-30m-rnnt-6000h" / "tokens.txt"
        if bpe_path.exists() and not tokens_path.exists():
            print("⚙️ Generating tokens.txt from bpe.model...")
            try:
                import sentencepiece as spm
                sp = spm.SentencePieceProcessor()
                sp.load(str(bpe_path))
                with open(tokens_path, "w", encoding="utf-8") as f:
                    for i in range(sp.get_piece_size()):
                        token = sp.id_to_piece(i)
                        f.write(f"{token} {i}\n")
                print("✅ Generated tokens.txt successfully!")
            except Exception as e:
                print(f"⚠️ Failed to generate tokens.txt: {e}")
    
    # 2. Zipformer 68M (Large RNN-T, ~250MB)
    download_model(
        repo_id="csukuangfj/sherpa-onnx-zipformer-vi-2025-04-20",
        local_dir="sherpa-onnx-zipformer-vi-2025-04-20",
        revision="0fc3fea3ccd9c50b439755fa8a6aba546cb3a7d4",
        allowed_patterns=["*.onnx", "tokens.txt", "bpe.model"]
    )
    
    # 3. ViBERT Punctuation (~440MB)
    if download_model(
        repo_id="welcomyou/vibert-capu-onnx",
        local_dir="vibert-capu",
        revision="a7754d037f4a9e29f7f3224f27acb60149eab874",
        allowed_patterns=["*.onnx", "*.txt", "*.json", "*.py", "vocabulary/*"]
    ):
        src = PROJECT_ROOT / "models" / "vibert-capu" / "verb-form-vocab.txt"
        dst = PROJECT_ROOT / "verb-form-vocab.txt"
        if src.exists():
            import shutil
            print(f"📋 Copying {src.name} to root folder...")
            shutil.copy2(src, dst)
            print("✅ Copied successfully!")

    # 4. Altunenes Speaker Segmentation ONNX (~6MB)
    download_model(
        repo_id="altunenes/speaker-diarization-community-1-onnx",
        local_dir="pyannote-onnx",
        revision="e2e09da94ae093a56cd5a60a09b138ae3da1959c",
        allowed_patterns=["segmentation-community-1.onnx"]
    )
    
    # 5. Pyannote Community-1 Embedding Split (~60MB)
    download_model(
        repo_id="welcomyou/pyannote-community-1-onnx-split",
        local_dir="pyannote-onnx",
        revision="cde44c2db938c8abb755853b9a87cb3179c47803",
        allowed_patterns=["embedding_encoder.onnx", "resnet_seg_1_weight.npy", "resnet_seg_1_bias.npy", "plda/*"]
    )

if __name__ == "__main__":
    main()