#!/bin/bash
# Script to download Piper TTS voices and place them in the correct structure
# Voices are downloaded from Hugging Face: https://huggingface.co/rhasspy/piper-voices

# Don't exit on error - we want to continue downloading other voices even if one fails
set +e

PIPER_DATA_DIR="./piper-data"
VOICES_BASE_URL="https://huggingface.co/rhasspy/piper-voices/resolve/main"

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${GREEN}Piper TTS Voice Downloader${NC}"
echo "================================"
echo ""

# Create piper-data directory if it doesn't exist
mkdir -p "$PIPER_DATA_DIR"

# Function to download a voice
download_voice() {
    local lang=$1
    local locale=$2
    local voice=$3
    local quality=$4
    
    local voice_name="${locale}-${voice}-${quality}"
    local voice_dir="${PIPER_DATA_DIR}/voices/${locale}/${voice}/${quality}"
    local onnx_file="${voice_dir}/${voice_name}.onnx"
    local json_file="${voice_dir}/${voice_name}.onnx.json"
    
    # Create directory structure
    mkdir -p "$voice_dir"
    
    # Check if voice already exists
    if [ -f "$onnx_file" ] && [ -f "$json_file" ]; then
        echo -e "${YELLOW}Voice ${voice_name} already exists, skipping...${NC}"
        return 0
    fi
    
    echo "Downloading ${voice_name}..."
    
    # Download .onnx.json file
    local json_url="${VOICES_BASE_URL}/${lang}/${locale}/${voice}/${quality}/${voice_name}.onnx.json"
    local curl_output=$(curl -L -f -w "\n%{http_code}" -o "$json_file" "$json_url" 2>&1)
    local http_code=$(echo "$curl_output" | tail -n1)
    local curl_error=$(echo "$curl_output" | head -n-1)
    
    if [ "$http_code" = "200" ] && [ -f "$json_file" ] && [ -s "$json_file" ]; then
        echo "  ✓ Downloaded ${voice_name}.onnx.json"
    else
        echo "  ✗ Failed to download ${voice_name}.onnx.json"
        echo "    URL: ${json_url}"
        echo "    HTTP Code: ${http_code:-unknown}"
        if [ -n "$curl_error" ]; then
            echo "    Error: $(echo "$curl_error" | head -n1)"
        fi
        echo "    This quality level may not be available for this voice."
        rm -f "$json_file"
        return 1
    fi
    
    # Download .onnx file
    local onnx_url="${VOICES_BASE_URL}/${lang}/${locale}/${voice}/${quality}/${voice_name}.onnx"
    curl_output=$(curl -L -f -w "\n%{http_code}" -o "$onnx_file" "$onnx_url" 2>&1)
    http_code=$(echo "$curl_output" | tail -n1)
    curl_error=$(echo "$curl_output" | head -n-1)
    
    if [ "$http_code" = "200" ] && [ -f "$onnx_file" ] && [ -s "$onnx_file" ]; then
        local file_size=$(stat -c%s "$onnx_file" 2>/dev/null || echo "0")
        local file_size_mb=$(echo "scale=2; $file_size / 1024 / 1024" | bc 2>/dev/null || echo "?")
        echo "  ✓ Downloaded ${voice_name}.onnx (${file_size_mb} MB)"
    else
        echo "  ✗ Failed to download ${voice_name}.onnx"
        echo "    URL: ${onnx_url}"
        echo "    HTTP Code: ${http_code:-unknown}"
        if [ -n "$curl_error" ]; then
            echo "    Error: $(echo "$curl_error" | head -n1)"
        fi
        echo "    This quality level may not be available for this voice."
        rm -f "$onnx_file" "$json_file"
        return 1
    fi
    
    echo -e "${GREEN}  ✓ Successfully downloaded ${voice_name}${NC}"
    return 0
}

# List of voices to download (based on the language detection function)
# Format: language_code locale voice quality
VOICES=(
    # English (US) - all quality levels
    "en en_US lessac low"
    "en en_US lessac medium"
    "en en_US lessac high"
    # English (GB)
    "en en_GB alba medium"
    
    # German
    "de de_DE thorsten medium"
    "de de_DE thorsten low"
    
    # French
    "fr fr_FR siwis medium"
    "fr fr_FR siwis low"
    
    # Spanish
    "es es_ES davefx medium"
    # Note: es_ES-davefx-low doesn't exist
    
    # Italian - riccardo doesn't exist, removing
    # "it it_IT riccardo medium" - not available
    # "it it_IT riccardo low" - not available
    
    # Russian
    "ru ru_RU ruslan medium"
    # Note: ru_RU-ruslan-low doesn't exist
    
    # Chinese
    "zh zh_CN huayan medium"
    
    # Arabic - hafez doesn't exist, removing
    # "ar ar_SA hafez medium" - not available
    
    # Polish
    "pl pl_PL darkman medium"
    
    # Portuguese - edresson doesn't exist, removing
    # "pt pt_BR edresson medium" - not available
    
    # Dutch
    "nl nl_NL mls medium"
    
    # Czech
    "cs cs_CZ jirka medium"
    
    # Turkish
    "tr tr_TR dfki medium"
    
    # Japanese - nanami doesn't exist, removing
    # "ja ja_JP nanami medium" - not available
    
    # Korean - kyungha doesn't exist, removing
    # "ko ko_KR kyungha medium" - not available
)

# Check if specific voices are requested
if [ $# -gt 0 ]; then
    VOICES=("$@")
fi

echo "Downloading ${#VOICES[@]} voice(s)..."
echo ""

SUCCESS=0
FAILED=0

for voice_spec in "${VOICES[@]}"; do
    # Parse voice specification
    read -r lang locale voice_name quality <<< "$voice_spec"
    
    # Validate voice specification
    if [ -z "$lang" ] || [ -z "$locale" ] || [ -z "$voice_name" ] || [ -z "$quality" ]; then
        echo -e "${YELLOW}⚠ Skipping invalid voice specification: '${voice_spec}'${NC}"
        ((FAILED++))
        echo ""
        continue
    fi
    
    if download_voice "$lang" "$locale" "$voice_name" "$quality"; then
        ((SUCCESS++))
    else
        ((FAILED++))
    fi
    echo ""
done

echo "================================"
echo -e "${GREEN}Download complete!${NC}"
echo "Successfully downloaded: $SUCCESS"
if [ $FAILED -gt 0 ]; then
    echo -e "${YELLOW}Failed: $FAILED${NC}"
fi
echo ""
echo "Voices are now in: $PIPER_DATA_DIR"
echo "This directory is mounted into the Docker container at /data"
echo ""
echo "To use these voices, restart your Docker containers:"
echo "  docker-compose down"
echo "  docker-compose up --build"