#!/bin/bash

# Setup script for Medical Image Analysis Tool
# This script handles dependency installation and environment setup

set -e  # Exit on any error

echo "=== Medical Image Analysis Tool Setup ==="
echo "Starting setup process at $(date)"

# Function to check if a Python package is installed
check_python_package() {
    python3 -c "import $1" 2>/dev/null && echo "✅ $1" || echo "❌ $1"
}

# Function to install Python packages with error handling
install_python_package() {
    local package=$1
    echo "Installing $package..."
    
    if pip install "$package" --no-cache-dir; then
        echo "✅ Successfully installed $package"
    else
        echo "❌ Failed to install $package"
        return 1
    fi
}

# Update pip first
echo "Updating pip..."
python3 -m pip install --upgrade pip

# Install core dependencies
echo "Installing core dependencies..."
install_python_package "torch>=2.5.1"
install_python_package "torchvision>=0.20.1"
install_python_package "transformers>=4.30.0"
install_python_package "gradio>=4.0.0"
install_python_package "numpy>=1.24.4"
install_python_package "pillow>=9.4.0"

# Install additional required packages
echo "Installing additional dependencies..."
install_python_package "albumentations"
install_python_package "einops"
install_python_package "opencv-python"
install_python_package "tqdm"
install_python_package "hydra-core"
install_python_package "omegaconf"

# Install SAM-2 if directory exists
if [ -d "segment-anything-2" ]; then
    echo "Installing SAM-2..."
    cd segment-anything-2
    pip install -e . --no-cache-dir
    cd ..
    echo "✅ SAM-2 installation completed"
else
    echo "⚠️  SAM-2 directory not found - cloning repository..."
    git clone https://github.com/facebookresearch/segment-anything-2.git
    cd segment-anything-2
    pip install -e . --no-cache-dir
    cd ..
    echo "✅ SAM-2 cloned and installed"
fi

# Verify installations
echo "Verifying installations..."
echo "Python packages status:"
check_python_package "torch"
check_python_package "torchvision"
check_python_package "transformers"
check_python_package "gradio"
check_python_package "albumentations"
check_python_package "einops"
check_python_package "cv2"
check_python_package "numpy"
check_python_package "PIL"

# Test torch and torchvision compatibility
echo "Testing PyTorch and Torchvision compatibility..."
python3 -c "
import torch
import torchvision
print(f'PyTorch version: {torch.__version__}')
print(f'Torchvision version: {torchvision.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print('✅ PyTorch and Torchvision are compatible')
" || echo "❌ PyTorch/Torchvision compatibility issue"

# Create necessary directories
echo "Creating necessary directories..."
mkdir -p logs
mkdir -p uploads
mkdir -p outputs

# Set permissions
chmod +x app.py

echo "=== Setup completed at $(date) ==="
echo "You can now run the application with: python3 app.py"

# Optional: Run a quick test
echo "Running quick test..."
python3 -c "
import sys
sys.path.append('.')
try:
    from app import check_dependencies
    deps = check_dependencies()
    print('Dependency check results:')
    for dep, status in deps.items():
        print(f'  {dep}: {\"✅\" if status else \"❌\"}')
except Exception as e:
    print(f'Test failed: {e}')
"