#!/bin/bash
# Quick setup script for VAD + Speaker Diarization

set -e

echo "=========================================="
echo "VAD + Speaker Diarization Setup"
echo "=========================================="

# Check if conda environment is active
if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
    echo "\n✓ Conda environment detected: $CONDA_DEFAULT_ENV"
    USE_CONDA=true
else
    echo "\n⚠️  No conda environment detected"
    USE_CONDA=false
fi

# Check Python version
echo -e "\n[1/6] Checking Python version..."
python_version=$(python --version 2>&1 | awk '{print $2}')
echo "Found Python $python_version"

if ! python -c "import sys; assert sys.version_info >= (3, 10)" 2>/dev/null; then
    echo "❌ Error: Python 3.10+ required"
    exit 1
fi
echo "✓ Python version OK"

# Check CUDA (optional)
echo -e "\n[2/6] Checking CUDA..."
if command -v nvidia-smi &> /dev/null; then
    cuda_version=$(nvidia-smi | grep "CUDA Version" | awk '{print $9}')
    echo "✓ CUDA $cuda_version detected"
    USE_CUDA=true
else
    echo "⚠️  No CUDA detected, will use CPU"
    USE_CUDA=false
fi

# Check FFmpeg
echo -e "\n[3/6] Checking FFmpeg..."
if command -v ffmpeg &> /dev/null; then
    echo "✓ FFmpeg installed"
else
    echo "⚠️  FFmpeg not found"
    echo "Install with: sudo apt-get install ffmpeg"
fi

# Setup environment
echo -e "\n[4/6] Setting up Python environment..."
if [ "$USE_CONDA" = true ]; then
    echo "✓ Using conda environment: $CONDA_DEFAULT_ENV"
else
    # Create virtual environment
    if [ ! -d "venv" ]; then
        python -m venv venv
        echo "✓ Virtual environment created"
    else
        echo "✓ Virtual environment already exists"
    fi
    # Activate virtual environment
    source venv/bin/activate
fi

# Install PyTorch
echo -e "\n[5/6] Installing PyTorch..."
if [ "$USE_CUDA" = true ]; then
    echo "Installing PyTorch with CUDA support..."
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
else
    echo "Installing PyTorch (CPU only)..."
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
fi

# Install dependencies
echo -e "\n[6/6] Installing dependencies..."
pip install -r requirements.txt

# Create directories
mkdir -p data outputs benchmarks

# Setup environment file
if [ ! -f ".env" ]; then
    cp .env.example .env
    echo "✓ Created .env file"
    echo "⚠️  Please edit .env and add your HF_TOKEN"
fi

echo -e "\n=========================================="
echo "✅ Setup complete!"
echo "=========================================="
echo -e "\nNext steps:"
if [ "$USE_CONDA" = true ]; then
    echo "1. Environment already active: $CONDA_DEFAULT_ENV ✓"
else
    echo "1. Activate environment: source venv/bin/activate"
fi
echo "2. Set HF token: export HF_TOKEN='your_token_here'"
echo "   Get token at: https://huggingface.co/settings/tokens"
echo "3. Accept model conditions at:"
echo "   https://huggingface.co/pyannote/speaker-diarization-3.1"
echo "4. Run demo: python vad_diarization.py"
echo "5. Run Gradio app: python app.py"
echo -e "\nFor more info, see README.md"
echo "=========================================="