#!/usr/bin/env bash

set -Eeuo pipefail

export JUPYTER_TOKEN="${JUPYTER_TOKEN:-huggingface}"
export JUPYTER_ROOT_DIR="${JUPYTER_ROOT_DIR:-/data/workspace}"
export BESZEL_DATA_DIR="${BESZEL_DATA_DIR:-/data/beszel-agent}"
export BESZEL_LISTEN="${BESZEL_LISTEN:-45876}"
export BESZEL_SYSTEM_NAME="${BESZEL_SYSTEM_NAME:-${SPACE_ID:-hf-space-jupyterlab}}"
export DISABLE_SSH="${DISABLE_SSH:-true}"
export DOCKER_HOST="${DOCKER_HOST:-}"
export SKIP_SYSTEMD="${SKIP_SYSTEMD:-true}"
export SKIP_GPU="${SKIP_GPU:-true}"

mkdir -p \
  "$JUPYTER_ROOT_DIR" \
  "$BESZEL_DATA_DIR" \
  "$HOME/.local/share/jupyter/runtime"

if [[ -z "${BESZEL_KEY:-}" || -z "${BESZEL_TOKEN:-}" || -z "${BESZEL_HUB_URL:-}" ]]; then
  echo "Missing required Space secrets: BESZEL_KEY, BESZEL_TOKEN, BESZEL_HUB_URL" >&2
  exit 1
fi

cleanup() {
  local exit_code=$?

  if [[ -n "${JUPYTER_PID:-}" ]]; then
    kill "$JUPYTER_PID" 2>/dev/null || true
  fi

  if [[ -n "${BESZEL_PID:-}" ]]; then
    kill "$BESZEL_PID" 2>/dev/null || true
  fi

  wait 2>/dev/null || true
  exit "$exit_code"
}

trap cleanup EXIT INT TERM

echo "Starting beszel-agent..."
KEY="$BESZEL_KEY" \
TOKEN="$BESZEL_TOKEN" \
HUB_URL="$BESZEL_HUB_URL" \
LISTEN="$BESZEL_LISTEN" \
DATA_DIR="$BESZEL_DATA_DIR" \
SYSTEM_NAME="$BESZEL_SYSTEM_NAME" \
DISABLE_SSH="$DISABLE_SSH" \
DOCKER_HOST="$DOCKER_HOST" \
SKIP_SYSTEMD="$SKIP_SYSTEMD" \
SKIP_GPU="$SKIP_GPU" \
beszel-agent &
BESZEL_PID=$!

echo "Starting JupyterLab on 0.0.0.0:7860..."
jupyter lab --config="$HOME/app/jupyter_server_config.py" &
JUPYTER_PID=$!

wait -n "$BESZEL_PID" "$JUPYTER_PID"
echo "A managed process exited, stopping the container." >&2