#!/bin/bash
# Helper script to download Piper TTS voice models

set -e

echo "=================================================="
echo "Piper TTS Voice Model Downloader"
echo "=================================================="
echo ""

# Create models directory
MODELS_DIR="$(dirname "$0")/piper_models"
mkdir -p "$MODELS_DIR"
echo "✓ Models directory: $MODELS_DIR"
echo ""

# Base URL for Piper voices
BASE_URL="https://huggingface.co/rhasspy/piper-voices/resolve/main"

# Function to download a voice
download_voice() {
    local voice_id=$1
    local lang_code=$2
    
    echo "Downloading $voice_id..."
    
    # Construct URLs
    local model_url="${BASE_URL}/${lang_code}/${lang_code}_${voice_id}.onnx"
    local config_url="${BASE_URL}/${lang_code}/${lang_code}_${voice_id}.onnx.json"
    
    local model_file="${MODELS_DIR}/${lang_code}-${voice_id}.onnx"
    local config_file="${MODELS_DIR}/${lang_code}-${voice_id}.onnx.json"
    
    # Check if already exists
    if [ -f "$model_file" ] && [ -f "$config_file" ]; then
        echo "  ✓ Already downloaded"
        return 0
    fi
    
    # Download with curl
    if command -v curl &> /dev/null; then
        echo "  Downloading model file..."
        curl -L -o "$model_file" "$model_url" 2>/dev/null || {
            echo "  ✗ Failed to download model"
            return 1
        }
        
        echo "  Downloading config file..."
        curl -L -o "$config_file" "$config_url" 2>/dev/null || {
            echo "  ✗ Failed to download config"
            return 1
        }
        
        # Verify files are not error pages
        if [ $(wc -c < "$model_file") -lt 1000 ]; then
            echo "  ✗ Download failed (file too small)"
            rm -f "$model_file" "$config_file"
            return 1
        fi
        
        echo "  ✓ Downloaded successfully"
        return 0
    else
        echo "  ✗ curl not found"
        return 1
    fi
}

echo "Select voices to download:"
echo ""
echo "Recommended starter pack (3 voices):"
echo "  1) en_US-lessac-medium (Male, American)"
echo "  2) en_US-amy-medium (Female, American)"
echo "  3) en_GB-alan-medium (Male, British)"
echo ""
echo "All English voices (5 voices):"
echo "  4) All of the above + 2 more"
echo ""
echo "Manual selection:"
echo "  5) Choose specific voices"
echo ""
echo "  0) Exit"
echo ""
read -p "Enter your choice (0-5): " choice

case $choice in
    1)
        echo ""
        echo "Downloading starter pack..."
        echo ""
        download_voice "lessac-medium" "en_US"
        download_voice "amy-medium" "en_US"
        download_voice "alan-medium" "en_GB"
        ;;
    2)
        echo ""
        echo "Downloading all English voices..."
        echo ""
        download_voice "lessac-medium" "en_US"
        download_voice "amy-medium" "en_US"
        download_voice "libritts-high" "en_US"
        download_voice "alan-medium" "en_GB"
        download_voice "alba-medium" "en_GB"
        ;;
    3)
        echo ""
        echo "Available voices:"
        echo "  en_US-lessac-medium (Male, American)"
        echo "  en_US-amy-medium (Female, American)"
        echo "  en_US-libritts-high (Neutral, American, High Quality)"
        echo "  en_GB-alan-medium (Male, British)"
        echo "  en_GB-alba-medium (Female, British)"
        echo ""
        echo "Enter voice IDs separated by spaces (e.g., lessac-medium amy-medium):"
        read -p "> " voices
        
        for voice in $voices; do
            if [[ $voice == *"GB"* ]]; then
                download_voice "$voice" "en_GB"
            else
                download_voice "$voice" "en_US"
            fi
        done
        ;;
    0)
        echo "Exiting..."
        exit 0
        ;;
    *)
        echo "Invalid choice"
        exit 1
        ;;
esac

echo ""
echo "=================================================="
echo "Download Summary"
echo "=================================================="
echo ""
echo "Downloaded files in: $MODELS_DIR"
ls -lh "$MODELS_DIR" 2>/dev/null || echo "No files found"
echo ""
echo "✓ Setup complete!"
echo ""
echo "Next steps:"
echo "  1. Test Piper: python test_piper_direct.py"
echo "  2. Start backend: uvicorn main:app --reload"
echo "  3. Select 'Piper TTS' in the UI"
echo ""