#!/bin/bash

# Print system information
echo "=== System Information ==="
echo "Python version: $(python3 --version)"
echo "CUDA version: $(nvcc --version 2>/dev/null || echo 'NVCC not found')"
echo "GPU information: $(nvidia-smi || echo 'NVIDIA-SMI not found')"
echo "==========================="

# Create writable directories in /tmp (where we have permissions)
echo "Creating writable directories..."
mkdir -p /tmp/comfyui_input
mkdir -p /tmp/comfyui_output
mkdir -p /tmp/comfyui_temp
mkdir -p /tmp/comfyui_user
mkdir -p /tmp/workflows

# Set proper permissions
chmod -R 777 /tmp/comfyui_input /tmp/comfyui_output /tmp/comfyui_temp /tmp/comfyui_user /tmp/workflows

# Download files from the dataset repository
echo "Downloading images from dataset repository..."
curl -L -o /tmp/comfyui_input/pose4.jpg "https://huggingface.co/datasets/Defter77/appdata/resolve/main/pose4.jpg" || echo "Failed to download pose4.jpg"
curl -L -o /tmp/comfyui_input/paita2.jpg "https://huggingface.co/datasets/Defter77/appdata/resolve/main/paita2.jpg" || echo "Failed to download paita2.jpg"

echo "Downloading workflow from dataset repository..."
curl -L -o /tmp/workflows/Workflow_12_11.json "https://huggingface.co/datasets/Defter77/appdata/resolve/main/Workflow_12_11.json" || echo "Failed to download Workflow_12_11.json"

# List downloaded files
echo "Verifying downloaded files:"
ls -la /tmp/comfyui_input/
ls -la /tmp/workflows/

# Install additional Python packages
echo "Installing required Python packages..."
pip install comfyui-frontend-package comfyui-workflow-templates || echo "Package installation failed, continuing anyway"

# Start ComfyUI in the background with custom directories
cd /app/ComfyUI
echo "Starting ComfyUI server..."
python3 main.py --listen 0.0.0.0 --port 8188 --input-directory /tmp/comfyui_input --output-directory /tmp/comfyui_output --temp-directory /tmp/comfyui_temp --user-dir /tmp/comfyui_user &
COMFY_PID=$!

# Wait for ComfyUI to start
echo "Waiting for ComfyUI server to initialize..."
sleep 30  # Give it more time to start

# Check if ComfyUI is running
if curl -s "http://localhost:8188/system_stats" > /dev/null; then
    echo "ComfyUI server is up and running!"
else
    echo "WARNING: ComfyUI server might not be running correctly."
    echo "Proceeding anyway..."
fi

# Start the web interface with updated workflow path
cd /app
echo "Starting Gradio interface..."
WORKFLOW_PATH="/tmp/workflows/Workflow_12_11.json" python3 app.py

# If the Gradio app exits, kill ComfyUI too
kill $COMFY_PID