#!/bin/bash

# =============================================================================
# AI Services Model Setup - Master Script
# =============================================================================
# This script sets up all AI models required for DeutschLernen:
#   - Vosk German speech recognition model (~500MB)
#   - Coqui TTS German model (~1.5GB)
# 
# Requirements:
#   - wget, unzip (for Vosk model)
#   - Python 3.8+ with pip (for Coqui TTS)
#   - ~2GB free disk space total
# 
# Usage:
#   ./setup-ai-models.sh [options]
# 
# Options:
#   --vosk-only       Download only Vosk model
#   --coqui-only      Download only Coqui TTS model
#   --help            Show this help message
#   --models-dir DIR  Custom models directory (default: ./models)
# 
# Example:
#   ./setup-ai-models.sh
#   ./setup-ai-models.sh --models-dir /opt/ai-models
# =============================================================================

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
NC='\033[0m' # No Color

# Default values
MODELS_DIR="./models"
VOSK_MODEL="vosk-model-de-0.22"
COQUI_MODEL="tts_models/de/deu/fairseq/vits"

# Parse arguments
VOSK_ONLY=false
COQUI_ONLY=false
SHOW_HELP=false

for arg in "$@"; do
    case "$arg" in
        --vosk-only)
            VOSK_ONLY=true
            ;;
        --coqui-only)
            COQUI_ONLY=true
            ;;
        --help)
            SHOW_HELP=true
            ;;
        --models-dir)
            MODELS_DIR="$2"
            shift
            ;;
        --models-dir=*)
            MODELS_DIR="${arg#*=}"
            ;;
    esac
done

# Function to print colored output
print_header() {
    echo -e "${PURPLE}==========================================${NC}"
    echo -e "${PURPLE}$1${NC}"
    echo -e "${PURPLE}==========================================${NC}"
    echo ""
}

print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Show help
if [ "$SHOW_HELP" = true ]; then
    echo "Usage: $0 [options]"
    echo ""
    echo "Options:"
    echo "  --vosk-only       Download only Vosk model"
    echo "  --coqui-only      Download only Coqui TTS model"
    echo "  --help            Show this help message"
    echo "  --models-dir DIR  Custom models directory (default: ./models)"
    echo ""
    echo "Examples:"
    echo "  $0                         # Download all models"
    echo "  $0 --vosk-only             # Download only Vosk model"
    echo "  $0 --coqui-only            # Download only Coqui TTS model"
    echo "  $0 --models-dir /opt/ai    # Use custom directory"
    echo ""
    exit 0
fi

print_header "DeutschLernen AI Services Setup"

# Check requirements
print_status "Checking system requirements..."

# Check wget for Vosk
if [ "$VOSK_ONLY" = false ] || [ "$COQUI_ONLY" = false ]; then
    if ! command -v wget &> /dev/null; then
        print_error "wget is required for Vosk model download"
        print_status "Install with: sudo apt-get install wget (Ubuntu/Debian)"
        exit 1
    fi
    print_success "wget: OK"
fi

# Check unzip for Vosk
if [ "$VOSK_ONLY" = false ] || [ "$COQUI_ONLY" = false ]; then
    if ! command -v unzip &> /dev/null; then
        print_error "unzip is required for Vosk model extraction"
        print_status "Install with: sudo apt-get install unzip (Ubuntu/Debian)"
        exit 1
    fi
    print_success "unzip: OK"
fi

# Check Python for Coqui
if [ "$COQUI_ONLY" = false ] || [ "$VOSK_ONLY" = false ]; then
    if ! command -v python3 &> /dev/null; then
        print_error "Python 3 is required for Coqui TTS"
        print_status "Install Python 3.8+ from https://www.python.org/downloads/"
        exit 1
    fi
    PYTHON_VERSION=$(python3 --version 2>&1 | awk '{print $2}')
    print_success "Python: ${PYTHON_VERSION}"
fi

# Create models directory
print_status "Creating models directory: ${MODELS_DIR}"
mkdir -p "${MODELS_DIR}"
print_success "Models directory created"
echo ""

# ============================================================================
# Vosk Speech Recognition Model
# ============================================================================
if [ "$COQUI_ONLY" = false ]; then
    print_header "Setting up Vosk German Model"
    print_status "Model: ${VOSK_MODEL} (~500MB)"
    
    VOSK_TARGET="${MODELS_DIR}/vosk"
    
    # Run Vosk setup script
    if [ -f "$(dirname "$0")/download-vosk-model.sh" ]; then
        print_status "Running Vosk setup script..."
        "$(dirname "$0")/download-vosk-model.sh" "${VOSK_TARGET}"
    else
        print_error "Vosk setup script not found. Running inline setup..."
        
        mkdir -p "${VOSK_TARGET}"
        cd "${VOSK_TARGET}"
        print_status "Downloading ${VOSK_MODEL}.zip..."
        wget "https://alphacephei.com/vosk/models/${VOSK_MODEL}.zip" -O "${VOSK_MODEL}.zip"
        print_status "Extracting..."
        unzip "${VOSK_MODEL}.zip"
        rm "${VOSK_MODEL}.zip"
        cd - > /dev/null
        
        print_success "Vosk model installed at: ${VOSK_TARGET}/${VOSK_MODEL}"
    fi
    
    print_success "Vosk model setup complete!"
    echo ""
fi

# ============================================================================
# Coqui TTS Model
# ============================================================================
if [ "$VOSK_ONLY" = false ]; then
    print_header "Setting up Coqui TTS German Model"
    print_status "Model: ${COQUI_MODEL} (~1.5GB)"
    
    COQUI_AUDIO="${MODELS_DIR}/tts-audio"
    
    # Run Coqui setup script
    if [ -f "$(dirname "$0")/download-coqui-model.sh" ]; then
        print_status "Running Coqui setup script..."
        "$(dirname "$0")/download-coqui-model.sh" "${COQUI_MODEL}" "${COQUI_AUDIO}"
    else
        print_error "Coqui setup script not found. Running inline setup..."
        
        print_status "Installing Coqui TTS..."
        pip install TTS
        
        mkdir -p "${COQUI_AUDIO}"
        
        print_status "Testing model loading (will auto-download if needed)..."
        python3 -c "from TTS.api import TTS; TTS(model_name='${COQUI_MODEL}'); print('OK')" || \
            print_warning "Model will download on first use"
        
        print_success "Coqui TTS setup complete!"
    fi
    
    print_success "Coqui TTS model setup complete!"
    echo ""
fi

# ============================================================================
# Summary
# ============================================================================
print_header "Setup Complete!"

echo "Models installed:"
if [ "$COQUI_ONLY" = false ]; then
    print_success "✓ Vosk: ${MODELS_DIR}/vosk/${VOSK_MODEL}"
fi
if [ "$VOSK_ONLY" = false ]; then
    print_success "✓ Coqui TTS: Auto-downloaded on first use"
fi

echo ""
print_status "Next steps:"
print_status "1. Install Python packages: pip install vosk TTS"
print_status "2. Update appsettings.json with model paths"
print_status "3. Start the application: dotnet run"
print_status "4. Test AI services via health endpoint: /health"
echo ""
print_status "For detailed instructions, see:"
print_status "  - scripts/ai-setup/download-vosk-model.sh"
print_status "  - scripts/ai-setup/download-coqui-model.sh"
print_status "  - docs/features/ai-services.md"
echo ""