#!/bin/bash

# MIMO Installation Script
# Compatible with conda environments and HuggingFace Spaces
#
# USAGE:
# 1. First activate conda environment: conda activate mimo
# 2. Then run this script: ./install.sh
#
# Requirements:
# - Python 3.10 (recommended via conda environment 'mimo')
# - CUDA-capable GPU (for local development)

set -e  # Exit on any error

echo "🚀 Starting MIMO installation..."

# Install system dependencies first
echo "🔧 Installing system dependencies..."
if [[ "$OSTYPE" == "darwin"* ]]; then
    # macOS - check for ffmpeg
    if ! command -v ffmpeg &> /dev/null; then
        echo "📦 Installing ffmpeg via Homebrew..."
        if command -v brew &> /dev/null; then
            brew install ffmpeg
        else
            echo "⚠️  Please install Homebrew first: https://brew.sh/"
            echo "    Then run: brew install ffmpeg"
            echo "    Or run: ./install_system_deps.sh"
        fi
    else
        echo "✅ ffmpeg already installed"
    fi
elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
    # Linux - attempt to install ffmpeg
    if ! command -v ffmpeg &> /dev/null; then
        echo "📦 Installing ffmpeg..."
        if command -v apt-get &> /dev/null; then
            sudo apt-get update && sudo apt-get install -y ffmpeg
        elif command -v yum &> /dev/null; then
            sudo yum install -y ffmpeg
        else
            echo "⚠️  Please install ffmpeg manually for your Linux distribution"
        fi
    else
        echo "✅ ffmpeg already installed"
    fi
fi

# Check if conda environment is activated
if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
    echo "🐍 Conda environment detected: $CONDA_DEFAULT_ENV"
    if [[ "$CONDA_DEFAULT_ENV" == "mimo" ]]; then
        echo "✅ Using MIMO conda environment"
        ENVIRONMENT="local"
        PYTHON_CMD="python"
        PIP_CMD="pip"
    else
        echo "⚠️  Warning: Expected 'mimo' conda environment, but found '$CONDA_DEFAULT_ENV'"
        ENVIRONMENT="local"
        PYTHON_CMD="python"
        PIP_CMD="pip"
    fi
elif [[ -n "$SPACE_ID" ]] || [[ -n "$HF_HOME" ]]; then
    echo "📦 Detected HuggingFace Spaces environment"
    ENVIRONMENT="hf"
    PYTHON_CMD="python3"
    PIP_CMD="pip3"
else
    echo "💻 Detected local development environment (non-conda)"
    ENVIRONMENT="local"
    PYTHON_CMD="python3"
    PIP_CMD="pip3"
fi

# Verify Python version
echo "🔍 Checking Python version..."
$PYTHON_CMD --version

# Fix typing-extensions first to resolve dependency conflicts
echo "🔧 Fixing typing-extensions version conflicts..."
$PIP_CMD install --upgrade "typing-extensions>=4.12.0"

# Install environment-specific PyTorch
if [[ "$ENVIRONMENT" == "local" ]]; then
    # For local development, we'll use versions from requirements_stable.txt
    # Skip separate PyTorch installation as it will be handled by requirements
    echo "⚡ PyTorch will be installed from requirements file for better compatibility"
else
    # HuggingFace Spaces - CPU optimized versions
    echo "☁️  Installing PyTorch for HuggingFace Spaces..."
    $PIP_CMD install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
fi

# Use stable requirements for better compatibility
echo "📚 Installing core dependencies..."

# First, completely clean any conflicting packages
echo "🧹 Cleaning existing installations to prevent conflicts..."
set +e
$PIP_CMD uninstall -y torch torchvision torchaudio tensorboard tensorflow accelerate diffusers transformers huggingface-hub 2>/dev/null || true
set -e

if [[ -f "requirements.txt" ]]; then
    echo "Using main requirements (MIMO-compatible versions)..."
    $PIP_CMD install -r requirements.txt
elif [[ -f "requirements_stable.txt" ]]; then
    echo "Using stable requirements..."
    $PIP_CMD install -r requirements_stable.txt
else
    echo "No requirements file found!"
    exit 1
fi

# Force MIMO-compatible HuggingFace ecosystem versions
echo "🔧 Ensuring MIMO-compatible HuggingFace ecosystem versions..."
$PIP_CMD install huggingface-hub==0.17.3 diffusers==0.23.1 transformers==4.35.2 accelerate==0.20.3 --force-reinstall

# Remove conflicting packages detection section since we already cleaned
# Everything will be handled by requirements_stable.txt

# Install additional dependencies based on environment
if [[ "$ENVIRONMENT" == "local" ]]; then
    # Local development - only install xformers if not Apple Silicon
    echo "� Installing optional performance enhancements..."
    set +e  # Allow this to fail gracefully

    # Skip xformers on Apple Silicon due to compilation issues
    if [[ $(uname -m) == "arm64" && $(uname -s) == "Darwin" ]]; then
        echo "⚠️  Skipping xformers on Apple Silicon (known compilation issues)"
        echo "   MIMO will work without xformers, just with slightly slower performance"
    else
        # Try to install xformers for other platforms
        $PIP_CMD install xformers==0.0.16
        if [ $? -eq 0 ]; then
            echo "✅ xformers installed successfully"
        else
            echo "⚠️  xformers installation failed, continuing without it"
            echo "   This is optional - MIMO will work fine without it"
        fi
    fi
    set -e

    # Skip TensorFlow to avoid conflicts - it's not essential for MIMO
    echo "📊 Installing TensorFlow (required for human segmentation)..."
    set +e
    $PIP_CMD install tensorflow==2.13.0
    if [ $? -eq 0 ]; then
        echo "✅ TensorFlow installed successfully"
    else
        echo "⚠️  TensorFlow installation failed - MIMO may not work fully without it"
    fi
    set -e
else
    # HuggingFace Spaces specific
    echo "🌐 Installing HuggingFace Spaces dependencies..."
    $PIP_CMD install spaces --upgrade
    # Ensure compatible gradio version for HF Spaces
    $PIP_CMD install gradio>=3.40.0 --upgrade
fi

# Verify installation
echo "✅ Verifying installation..."
$PYTHON_CMD -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"
$PYTHON_CMD -c "import transformers; print(f'Transformers version: {transformers.__version__}')"
$PYTHON_CMD -c "import diffusers; print(f'Diffusers version: {diffusers.__version__}')"
$PYTHON_CMD -c "import gradio; print(f'Gradio version: {gradio.__version__}')"
$PYTHON_CMD -c "import accelerate; print(f'Accelerate version: {accelerate.__version__}')"

# Optional: Check xformers
set +e
$PYTHON_CMD -c "
try:
    import xformers
    print(f'xformers version: {xformers.__version__}')
except ImportError:
    print('xformers: Not installed (optional)')
"
set -e

echo "🎉 Installation completed successfully!"
echo "Environment: $ENVIRONMENT"

# Optional: Download MIMO model
echo ""
echo "📥 Do you want to download the MIMO model from ModelScope? (y/n)"
read -r download_model
if [[ "$download_model" =~ ^[Yy]$ ]]; then
    echo "🔽 Downloading MIMO model..."
    $PYTHON_CMD -c "
from modelscope import snapshot_download
import os
print('Downloading MIMO model from ModelScope...')
model_dir = snapshot_download(model_id='iic/MIMO', cache_dir='./pretrained_weights')
print(f'Model downloaded successfully to: {model_dir}')
"
    if [ $? -eq 0 ]; then
        echo "✅ Model download completed!"
    else
        echo "⚠️  Model download failed. You can download it later using:"
        echo "   python -c \"from modelscope import snapshot_download; snapshot_download(model_id='iic/MIMO', cache_dir='./pretrained_weights')\""
    fi
else
    echo "⏭️  Skipping model download. You can download it later using:"
    echo "   python -c \"from modelscope import snapshot_download; snapshot_download(model_id='iic/MIMO', cache_dir='./pretrained_weights')\""
fi