#!/bin/bash

# ZenVision AI Subtitle Generator - Installation Script
# Instala todas las dependencias y modelos necesarios

set -e

echo "🚀 Instalando ZenVision AI Subtitle Generator..."
echo "=================================================="

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

# Function to print colored output
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"
}

# Check Python version
print_status "Verificando versión de Python..."
python_version=$(python3 --version 2>&1 | awk '{print $2}' | cut -d. -f1,2)
required_version="3.8"

if [ "$(printf '%s\n' "$required_version" "$python_version" | sort -V | head -n1)" = "$required_version" ]; then
    print_success "Python $python_version encontrado"
else
    print_error "Python $required_version o superior requerido. Encontrado: $python_version"
    exit 1
fi

# Check if running on supported OS
print_status "Detectando sistema operativo..."
if [[ "$OSTYPE" == "linux-gnu"* ]]; then
    OS="linux"
    print_success "Sistema Linux detectado"
elif [[ "$OSTYPE" == "darwin"* ]]; then
    OS="macos"
    print_success "Sistema macOS detectado"
elif [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "cygwin" ]]; then
    OS="windows"
    print_success "Sistema Windows detectado"
else
    print_warning "Sistema operativo no reconocido: $OSTYPE"
    OS="unknown"
fi

# Install system dependencies
print_status "Instalando dependencias del sistema..."

if [[ "$OS" == "linux" ]]; then
    if command -v apt-get &> /dev/null; then
        sudo apt-get update
        sudo apt-get install -y ffmpeg python3-dev python3-pip build-essential
        print_success "Dependencias de Ubuntu/Debian instaladas"
    elif command -v yum &> /dev/null; then
        sudo yum install -y ffmpeg python3-devel python3-pip gcc gcc-c++
        print_success "Dependencias de CentOS/RHEL instaladas"
    elif command -v pacman &> /dev/null; then
        sudo pacman -S --noconfirm ffmpeg python python-pip base-devel
        print_success "Dependencias de Arch Linux instaladas"
    else
        print_warning "Gestor de paquetes no reconocido. Instala manualmente: ffmpeg, python3-dev"
    fi
elif [[ "$OS" == "macos" ]]; then
    if command -v brew &> /dev/null; then
        brew install ffmpeg
        print_success "FFmpeg instalado via Homebrew"
    else
        print_warning "Homebrew no encontrado. Instala FFmpeg manualmente"
    fi
fi

# Create virtual environment
print_status "Creando entorno virtual..."
if [ ! -d "venv" ]; then
    python3 -m venv venv
    print_success "Entorno virtual creado"
else
    print_warning "Entorno virtual ya existe"
fi

# Activate virtual environment
print_status "Activando entorno virtual..."
source venv/bin/activate
print_success "Entorno virtual activado"

# Upgrade pip
print_status "Actualizando pip..."
pip install --upgrade pip setuptools wheel

# Install PyTorch with CUDA support if available
print_status "Detectando soporte CUDA..."
if command -v nvidia-smi &> /dev/null; then
    print_success "NVIDIA GPU detectada, instalando PyTorch con CUDA"
    pip install torch torchaudio --index-url https://download.pytorch.org/whl/cu118
else
    print_warning "No se detectó GPU NVIDIA, instalando PyTorch CPU"
    pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu
fi

# Install Python dependencies
print_status "Instalando dependencias de Python..."
pip install -r requirements.txt
print_success "Dependencias de Python instaladas"

# Download spaCy models
print_status "Descargando modelos de spaCy..."
python -m spacy download en_core_web_sm
python -m spacy download es_core_news_sm

# Try to download additional language models
for lang in fr de it pt; do
    print_status "Intentando descargar modelo de spaCy para $lang..."
    python -m spacy download ${lang}_core_news_sm || print_warning "Modelo $lang no disponible"
done

print_success "Modelos de spaCy descargados"

# Download NLTK data
print_status "Descargando datos de NLTK..."
python -c "
import nltk
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('vader_lexicon')
print('Datos de NLTK descargados')
"

# Create necessary directories
print_status "Creando directorios necesarios..."
mkdir -p ~/.zenvision/cache
mkdir -p ~/.zenvision/models
mkdir -p /tmp/zenvision
print_success "Directorios creados"

# Test installation
print_status "Probando instalación..."
python -c "
import torch
import whisper
import transformers
import gradio
import moviepy
import librosa
import cv2
import spacy
print('✅ Todas las librerías principales importadas correctamente')

# Test CUDA availability
if torch.cuda.is_available():
    print(f'✅ CUDA disponible: {torch.cuda.get_device_name(0)}')
else:
    print('⚠️ CUDA no disponible, usando CPU')

# Test Whisper
try:
    model = whisper.load_model('tiny')
    print('✅ Whisper cargado correctamente')
except Exception as e:
    print(f'❌ Error cargando Whisper: {e}')
"

# Create launcher script
print_status "Creando script de lanzamiento..."
cat > run_zenvision.sh << 'EOF'
#!/bin/bash
cd "$(dirname "$0")"
source venv/bin/activate
python app.py
EOF

chmod +x run_zenvision.sh
print_success "Script de lanzamiento creado: ./run_zenvision.sh"

# Display system information
print_status "Información del sistema:"
echo "  - Python: $(python --version)"
echo "  - PyTorch: $(python -c 'import torch; print(torch.__version__)')"
echo "  - CUDA disponible: $(python -c 'import torch; print(torch.cuda.is_available())')"
echo "  - Dispositivos CUDA: $(python -c 'import torch; print(torch.cuda.device_count())')"

# Calculate approximate model size
print_status "Calculando tamaño aproximado de modelos..."
echo "  - Whisper large-v2: ~1.5 GB"
echo "  - BERT multilingual: ~400 MB"
echo "  - RoBERTa sentiment: ~200 MB"
echo "  - DistilRoBERTa emotion: ~300 MB"
echo "  - Modelos de traducción: ~500 MB"
echo "  - Otros modelos: ~300 MB"
echo "  - TOTAL APROXIMADO: ~3.2 GB"

print_success "¡Instalación completada!"
echo ""
echo "=================================================="
echo "🎬 ZenVision AI Subtitle Generator está listo!"
echo "=================================================="
echo ""
echo "Para ejecutar la aplicación:"
echo "  1. Activar entorno virtual: source venv/bin/activate"
echo "  2. Ejecutar aplicación: python app.py"
echo "  3. O usar el script: ./run_zenvision.sh"
echo ""
echo "La aplicación estará disponible en: http://localhost:7860"
echo ""
echo "Para más información, consulta el README.md"
echo ""
print_success "¡Disfruta usando ZenVision! 🚀"