#!/bin/bash
set -e

# Global PIDs for cleanup
VNC_PID=""
FLUXBOX_PID=""

# Cleanup function to terminate background processes on script exit
cleanup() {
    echo "Cleaning up background processes..."
    # Kill processes in reverse order of startup. The '|| true' prevents errors if a process is already dead.
    if [ -n "$FLUXBOX_PID" ]; then kill -TERM $FLUXBOX_PID 2>/dev/null || true; fi
    if [ -n "$VNC_PID" ]; then kill -TERM $VNC_PID 2>/dev/null || true; fi
    echo "Cleanup complete."
}
trap cleanup EXIT

# Xvfb is now started by xvfb-run in the Dockerfile ENTRYPOINT.
# The DISPLAY variable will be set automatically by xvfb-run.


# It's safer to source conda.sh directly
source /opt/conda/etc/profile.d/conda.sh
conda activate camo

# Ensure the persistent data directory exists before we try to use it for the lock file.
mkdir -p /app/persistent-data

# --- One-time Initialization ---
# On first launch, multiple instances starting at once can cause a race condition
# during the download/extraction of the Camoufox distribution. To prevent this,
# we run a single dummy instance first, wait for it to become healthy (which
# indicates setup is complete), and then kill it. A lock file ensures this
# only happens on the very first start of the container.
INIT_LOCK_FILE="/app/persistent-data/camoufox.initialized"
if [ ! -f "$INIT_LOCK_FILE" ]; then
    echo "First start detected. Performing one-time Camoufox initialization..."

    # Start a single dummy instance in the background, logging to a file.
    # It will perform the necessary downloads and setup.
    INIT_LOG="/tmp/camoufox_init.log"
    rm -f "$INIT_LOG" # Ensure log file is clean before starting
    python3 -u camoufox_server.py --port 9999 --num-instances 1 > "$INIT_LOG" 2>&1 &
    INIT_PID=$!

    # Wait for the server to log that it's started, which is a reliable signal
    # that all one-time downloads and setup tasks are complete.
    echo "Waiting for initialization to complete (max 120s)..."
    
    end_time=$((SECONDS + 120))
    INIT_SUCCESS=false
    while [ $SECONDS -lt $end_time ]; do
        # The camoufox library logs "Websocket endpoint:" when it's ready.
        # This is a more reliable signal than a custom log message from our script.
        if grep -q "Websocket endpoint: ws://0.0.0.0:9999" "$INIT_LOG"; then
            INIT_SUCCESS=true
            break
        fi
        # Also check if the initialization process died unexpectedly
        if ! ps -p $INIT_PID > /dev/null; then
            echo "Initialization process died unexpectedly."
            break
        fi
        sleep 2
    done

    if [ "$INIT_SUCCESS" = true ]; then
        echo "Initialization successful."
    else
        echo "Initialization timed out or failed. The main server might fail to start."
        echo "--- Initialization Log ---"
        cat "$INIT_LOG"
        echo "--------------------------"
    fi

    # Cleanly terminate the dummy server.
    echo "Shutting down initialization server..."
    kill -TERM $INIT_PID
    wait $INIT_PID 2>/dev/null || true # Wait for it to exit, ignore error code

    # Create the lock file to prevent this from running again.
    touch "$INIT_LOCK_FILE"
    echo "Initialization complete. Proceeding with normal startup."
else
    echo "Initialization already complete. Skipping."
fi
# --- End Initialization ---

# Start supporting services (VNC, window manager)
echo "Starting VNC server on port 5900..."
# The -noxdamage flag is added to improve compatibility with VNC clients like the one on macOS.
# The '-localhost no' part was likely a typo and has been removed as the default is to allow non-localhost connections.
x11vnc -forever -usepw -display $DISPLAY -rfbport 5900 -o /var/log/x11vnc.log -shared -noxdamage &
VNC_PID=$!

echo "Starting Fluxbox window manager..."
fluxbox > /var/log/fluxbox.log 2>&1 &
FLUXBOX_PID=$!

# Start main application
echo "Starting Camoufox server with arguments: $@"
exec python3 -u camoufox_server.py "$@"