#!/bin/bash

# Start Desktop Environment Script
# This script initializes VNC server, noVNC, and the Gradio application

set -e

echo "=========================================="
echo "Starting X11 Desktop Environment"
echo "=========================================="

# Function to cleanup on exit
cleanup() {
    echo "Cleaning up..."
    pkill -u vncuser Xtigervnc || true
    pkill -u vncuser websockify || true
    pkill -u vncuser python3 || true
}
trap cleanup EXIT INT TERM

# Set display resolution (can be customized)
export RESOLUTION=${RESOLUTION:-1920x1080}
export DEPTH=${DEPTH:-24}

# Start VNC server as vncuser
echo "Starting VNC server on display ${DISPLAY}..."
su - vncuser -c "vncserver ${DISPLAY} -geometry ${RESOLUTION} -depth ${DEPTH} -localhost no -SecurityTypes None" || {
    echo "VNC server failed to start, trying to clean existing sessions..."
    su - vncuser -c "vncserver -kill ${DISPLAY}" || true
    sleep 2
    su - vncuser -c "vncserver ${DISPLAY} -geometry ${RESOLUTION} -depth ${DEPTH} -localhost no -SecurityTypes None"
}

# Wait for VNC server to be ready
echo "Waiting for VNC server to be ready..."
sleep 3

# Start noVNC websocket proxy with WSS support
echo "Starting noVNC WSS on port ${NO_VNC_PORT}..."
websockify --web=/usr/share/novnc --cert=/etc/ssl/certs/selfsigned.crt --key=/etc/ssl/private/selfsigned.key ${NO_VNC_PORT} localhost:${VNC_PORT} &
NOVNC_PID=$!

# Wait for noVNC to be ready
sleep 2

# Start FastAPI agent service
echo "Starting Agent API on port 8000..."
cd /app
python3 -m uvicorn agent.api:app --host 0.0.0.0 --port 8000 &
API_PID=$!

# Wait for API to be ready
sleep 2

# Start Gradio application
echo "Starting Gradio interface on port 7860..."
python3 app.py &
GRADIO_PID=$!

echo "=========================================="
echo "Services started successfully!"
echo "=========================================="
echo "noVNC URL: http://localhost:${NO_VNC_PORT}/vnc.html"
echo "Gradio UI: http://localhost:7860"
echo "Agent API: http://localhost:8000/docs"
echo "=========================================="

# Keep container running and monitor services
while true; do
    # Check if services are still running
    if ! kill -0 $NOVNC_PID 2>/dev/null; then
        echo "noVNC died, restarting..."
        websockify --web=/usr/share/novnc ${NO_VNC_PORT} localhost:${VNC_PORT} &
        NOVNC_PID=$!
    fi
    
    if ! kill -0 $API_PID 2>/dev/null; then
        echo "Agent API died, restarting..."
        python3 -m uvicorn agent.api:app --host 0.0.0.0 --port 8000 &
        API_PID=$!
    fi
    
    if ! kill -0 $GRADIO_PID 2>/dev/null; then
        echo "Gradio died, restarting..."
        python3 app.py &
        GRADIO_PID=$!
    fi
    
    sleep 10
done