#!/bin/bash
set -e

echo "🚀 Checking system configuration for Dataset Enrichment App..."

# Remove any old overrides so we start fresh
rm -f docker-compose.override.yml

# 1. Check for NVIDIA GPU
if command -v nvidia-smi &> /dev/null; then
    echo "✅ NVIDIA GPU detected."
    
    # 2. Check if the NVIDIA Container Toolkit is installed
    if ! command -v nvidia-ctk &> /dev/null; then
        echo "⚠️ NVIDIA Container Toolkit is missing."
        echo "   Without it, Docker cannot access your GPU, and the app will run in CPU-only mode."
        read -p "Would you like to install it automatically now? (requires sudo) [y/N] " install_choice
        if [[ "$install_choice" =~ ^[Yy]$ ]]; then
            echo "Installing NVIDIA Toolkit..."
            curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
            curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
                sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
                sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
            sudo apt-get update
            sudo apt-get install -y nvidia-container-toolkit
            sudo nvidia-ctk runtime configure --runtime=docker
            sudo systemctl restart docker
            echo "✅ NVIDIA Toolkit installed successfully!"
        else
            echo "⚠️ Skipping installation. Proceeding without GPU support."
        fi
    fi

    # 3. If Toolkit is available, configure CDI and enable GPU passthrough
    if command -v nvidia-ctk &> /dev/null; then
        if [ ! -f "/etc/cdi/nvidia.yaml" ]; then
            echo "⚙️ Generating Docker CDI configuration..."
            sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml
        fi
        
        echo "🔌 Enabling GPU support via docker-compose.override.yml..."
        cat <<EOF > docker-compose.override.yml
services:
  backend:
    devices:
      - nvidia.com/gpu=all
EOF
    fi
else
    echo "ℹ️ No NVIDIA GPU detected. Running in CPU-only mode."
fi

if docker info &> /dev/null; then
    echo "📦 Starting containers..."
    docker compose up -d --build
    PORT=${WEB_PORT:-8080}
else
    echo "⚠️ Docker is not accessible (or requires root permissions)."
    echo "🚀 Starting app in local development mode..."
    uv run uvicorn backend.main:app --host 0.0.0.0 --port 8000 &
    npm --prefix frontend run dev -- --host 0.0.0.0 &
    PORT=5173
fi

# Get local IP for convenience
LOCAL_IP=$(hostname -I | awk '{print $1}' || echo "localhost")

echo ""
echo "=========================================================="
echo "✅ App is running!"
echo "🌐 Access it locally at: http://localhost:$PORT"
echo "📱 Access it on your network at: http://$LOCAL_IP:$PORT"
echo "=========================================================="