#!/bin/bash
set -euo pipefail

APP_DIR="${APP_DIR:-/app}"
APP_FILE="${APP_FILE:-${APP_DIR}/app.py}"
PREFLIGHT_FILE="${PREFLIGHT_FILE:-${APP_DIR}/runtime_preflight.py}"
LOG_FILE="${LOG_FILE:-/workspace/gui.log}"
PORT="${GRADIO_SERVER_PORT:-7860}"
STOP_TIMEOUT="${STOP_TIMEOUT:-30}"

if [ ! -f "$APP_FILE" ]; then
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    APP_DIR="$SCRIPT_DIR"
    APP_FILE="${APP_DIR}/app.py"
fi

if [ ! -f "$APP_FILE" ]; then
    echo "ERROR: app.py not found. Set APP_DIR or APP_FILE."
    exit 1
fi

mkdir -p "$(dirname "$LOG_FILE")"

echo "Stopping existing Gradio/app.py process..."
mapfile -t PIDS < <(pgrep -f "python.*app.py|${APP_FILE}" || true)

if [ "${#PIDS[@]}" -gt 0 ]; then
    for pid in "${PIDS[@]}"; do
        if [ "$pid" = "$$" ]; then
            continue
        fi
        if [ "$pid" = "1" ] && [ "${FORCE_PID1_RESTART:-0}" != "1" ]; then
            echo "ERROR: app.py is PID 1. Killing it will stop this container before the script can restart it."
            echo "Restart the pod/container, or rerun with FORCE_PID1_RESTART=1 if your runtime restarts PID 1 automatically."
            exit 2
        fi
        kill -TERM "$pid" 2>/dev/null || true
    done

    deadline=$((SECONDS + STOP_TIMEOUT))
    while [ "$SECONDS" -lt "$deadline" ]; do
        alive=0
        for pid in "${PIDS[@]}"; do
            if [ "$pid" != "$$" ] && kill -0 "$pid" 2>/dev/null; then
                alive=1
                break
            fi
        done
        [ "$alive" -eq 0 ] && break
        sleep 1
    done

    for pid in "${PIDS[@]}"; do
        if [ "$pid" != "$$" ] && kill -0 "$pid" 2>/dev/null; then
            echo "Force-stopping PID $pid..."
            kill -KILL "$pid" 2>/dev/null || true
        fi
    done
else
    echo "No existing app.py process found."
fi

echo "Clearing Python/CUDA process state by starting a fresh app process..."
if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi || true
fi

if [ -f "$PREFLIGHT_FILE" ]; then
    echo "Running runtime preflight..."
    python "$PREFLIGHT_FILE"
else
    echo "Runtime preflight not found at ${PREFLIGHT_FILE}; continuing."
fi

cd "$APP_DIR"
echo "Starting Gradio on port ${PORT}..."
nohup python "$APP_FILE" > "$LOG_FILE" 2>&1 &
NEW_PID=$!
disown "$NEW_PID" 2>/dev/null || true

echo "Started PID ${NEW_PID}. Logs: ${LOG_FILE}"

if command -v curl >/dev/null 2>&1; then
    echo "Waiting for GUI to respond..."
    for _ in $(seq 1 90); do
        if curl -fsS "http://127.0.0.1:${PORT}/" >/dev/null 2>&1; then
            echo "GUI is ready: http://127.0.0.1:${PORT}/"
            exit 0
        fi
        if ! kill -0 "$NEW_PID" 2>/dev/null; then
            echo "ERROR: app process exited during startup. Last log lines:"
            tail -n 80 "$LOG_FILE" || true
            exit 1
        fi
        sleep 2
    done
    echo "Started, but GUI did not respond within timeout. Check logs: ${LOG_FILE}"
fi