#!/bin/bash
# Refacer Docker setup script
# Run this script from within the refacer/docker folder.

# --------------------------------------------
# Step 1: Download the Model if not available
# --------------------------------------------
MODEL_URL="https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
MODEL_PATH="../inswapper_128.onnx"
if [ ! -f "$MODEL_PATH" ]; then
    echo "Downloading inswapper_128.onnx model..."
    wget -O "$MODEL_PATH" "$MODEL_URL"
    if [ $? -ne 0 ]; then
        echo "Error: Failed to download the model. Please check your internet connection."
        exit 1
    fi
    echo "Model downloaded successfully!"
else
    echo "Model already exists at $MODEL_PATH."
fi

# --------------------------------------------
# Step 2: Stop Existing Docker Containers
# --------------------------------------------
echo "Stopping any running 'refacer' containers..."
docker stop -t 0 refacer 2>/dev/null || echo "No running container found."

# --------------------------------------------
# Step 3: Build the Docker Image
# --------------------------------------------
echo "Building the Docker image 'refacer' using Dockerfile.nvidia..."
docker build -t refacer -f Dockerfile.nvidia .
if [ $? -ne 0 ]; then
    echo "Error: Docker image build failed."
    exit 1
fi
echo "Docker image built successfully."

# --------------------------------------------
# Step 4: Run the Docker Container
# --------------------------------------------
echo "Starting the Docker container with NVIDIA GPU support..."
docker run --rm --name refacer \
    -v $(pwd)/..:/refacer \
    -p 7860:7860 \
    --gpus all \
    refacer python3 app.py --server_name 0.0.0.0 &

# Wait for the server to start
echo "Waiting for the Refacer server to start..."
sleep 5

# --------------------------------------------
# Step 5: Open Browser Automatically
# --------------------------------------------
echo "Opening Refacer UI in your browser..."
if command -v google-chrome >/dev/null 2>&1; then
    google-chrome --new-window "http://127.0.0.1:7860" &
elif command -v xdg-open >/dev/null 2>&1; then
    xdg-open "http://127.0.0.1:7860" &
elif command -v open >/dev/null 2>&1; then
    open "http://127.0.0.1:7860" &
else
    echo "Please open your browser and navigate to: http://127.0.0.1:7860"
fi