#!/usr/bin/env bash
set -Eeuo pipefail

echo "Starting LC0 Cloud bootstrap..."

ensure_tailscale() {
  if command -v tailscaled >/dev/null 2>&1 && command -v tailscale >/dev/null 2>&1; then
    return 0
  fi

  if ! command -v apt-get >/dev/null 2>&1; then
    echo "tailscale is not installed and apt-get is unavailable; install tailscale manually."
    exit 1
  fi

  echo "Installing tailscale runtime..."
  export DEBIAN_FRONTEND=noninteractive

  apt-get update
  apt-get install -y --no-install-recommends ca-certificates curl gnupg iproute2

  . /etc/os-release
  CODENAME="${VERSION_CODENAME:-}"
  if [[ -z "$CODENAME" ]]; then
    echo "Unable to determine OS codename for tailscale package setup."
    exit 1
  fi

  install -d -m 0755 /usr/share/keyrings
  curl -fsSL "https://pkgs.tailscale.com/stable/ubuntu/${CODENAME}.noarmor.gpg" \
    | tee /usr/share/keyrings/tailscale-archive-keyring.gpg >/dev/null
  curl -fsSL "https://pkgs.tailscale.com/stable/ubuntu/${CODENAME}.tailscale-keyring.list" \
    | tee /etc/apt/sources.list.d/tailscale.list >/dev/null

  apt-get update
  apt-get install -y --no-install-recommends tailscale
}

WORKDIR="${WORKDIR:-/workspace}"
LOG_DIR="${LOG_DIR:-$WORKDIR/logs}"
TS_STATE="${TS_STATE:-$WORKDIR/tailscaled.state}"
TS_SOCKET="${TS_SOCKET:-$WORKDIR/tailscaled.sock}"
TS_KEY_FILE="${TS_KEY_FILE:-$WORKDIR/.ts_authkey}"
BRIDGE_BIN="${BRIDGE_BIN:-$WORKDIR/uci-bridge-linux}"
LC0_BIN="${LC0_BIN:-$WORKDIR/lc0/build/lc0}"
WEIGHTS_FILE="${WEIGHTS_FILE:-$WORKDIR/BT4-332.pb}"

LISTEN_ADDR="${LISTEN_ADDR:-0.0.0.0:5555}"
MGMT_ADDR="${MGMT_ADDR:-0.0.0.0:53350}"
TS_HOSTNAME="${TS_HOSTNAME:-aichessgm-vast-5090}"
TS_ACCEPT_ROUTES="${TS_ACCEPT_ROUTES:-0}"
LISTEN_PORT="${LISTEN_ADDR##*:}"
MGMT_PORT="${MGMT_ADDR##*:}"

ensure_tailscale

mkdir -p "$LOG_DIR"

test -x "$BRIDGE_BIN" || { echo "Missing executable: $BRIDGE_BIN"; exit 1; }
test -x "$LC0_BIN" || { echo "Missing executable: $LC0_BIN"; exit 1; }
test -f "$WEIGHTS_FILE" || { echo "Missing weights: $WEIGHTS_FILE"; exit 1; }

if pgrep -f "$BRIDGE_BIN" >/dev/null 2>&1 && ss -ltn | grep -q ":${LISTEN_PORT}[[:space:]]"; then
  echo "Bridge already running on port ${LISTEN_PORT}. Exiting."
  exit 0
fi

TS_AUTHKEY="${TS_AUTHKEY:-}"
if [[ -z "$TS_AUTHKEY" && -f "$TS_KEY_FILE" ]]; then
  TS_AUTHKEY="$(cat "$TS_KEY_FILE")"
fi

if ! pgrep -x tailscaled >/dev/null 2>&1; then
  echo "Starting tailscaled..."
  nohup tailscaled \
    --tun=userspace-networking \
    --state="$TS_STATE" \
    --socket="$TS_SOCKET" \
    >"$LOG_DIR/tailscaled.log" 2>&1 &
fi

for _ in {1..40}; do
  [[ -S "$TS_SOCKET" ]] && break
  sleep 0.25
done

if ! pgrep -f "$BRIDGE_BIN" >/dev/null 2>&1; then
  echo "Starting UCI bridge on $LISTEN_ADDR ..."
  nohup "$BRIDGE_BIN" \
    --listen "$LISTEN_ADDR" \
    --mgmt-listen "$MGMT_ADDR" \
    --lc0 "$LC0_BIN" \
    --weights "$WEIGHTS_FILE" \
    --backend cuda-fp16 \
    --engine-name lc0-cloud \
    >"$LOG_DIR/bridge.log" 2>&1 &
fi

echo "Ensuring Tailscale session..."
if ! tailscale --socket="$TS_SOCKET" ip -4 >/dev/null 2>&1; then
  echo "Tailscale not ready yet; attempting up in background..."
  if ! pgrep -f "tailscale.*--socket=$TS_SOCKET.* up" >/dev/null 2>&1; then
    (
      TS_UP_ARGS=(--hostname "$TS_HOSTNAME")
      [[ -n "${TS_AUTHKEY:-}" ]] && TS_UP_ARGS+=(--auth-key "$TS_AUTHKEY")
      [[ "${TS_ACCEPT_ROUTES:-0}" == "1" ]] && TS_UP_ARGS+=(--accept-routes)

      if command -v timeout >/dev/null 2>&1; then
        timeout 15s tailscale --socket="$TS_SOCKET" up "${TS_UP_ARGS[@]}" || true
      else
        tailscale --socket="$TS_SOCKET" up "${TS_UP_ARGS[@]}" || true
      fi
    ) >>"$LOG_DIR/tailscale-up.log" 2>&1 &
  fi
fi

sleep 2

echo "LC0 Cloud bootstrap complete."
tailscale --socket="$TS_SOCKET" ip -4 || true
ss -ltnp | egrep "${LISTEN_PORT}|${MGMT_PORT}" || true
