#!/usr/bin/env bash
# willow-serve — turn willow-mcp OAuth serve mode on/off at will.
#
# Manages a systemd --user unit for the `--serve` HTTP process AND toggles the
# matching http entry in this repo's .mcp.json, so the client connects only when
# serve is on. No hand-editing of config, no code change per toggle.
#
#   willow-serve install   # one-time: write + load the systemd user unit
#   willow-serve on         # start serve + add the .mcp.json entry   (then /mcp)
#   willow-serve off        # stop serve  + remove the .mcp.json entry (then /mcp)
#   willow-serve status     # unit state + whether the entry is present
#   willow-serve logs        # follow the serve process logs (journalctl)
#
# Kart worker (kartikeya) — `willow-mcp worker` as a lane-instanced user unit:
#   willow-serve worker-install          # one-time: write + load willow-mcp-worker@.service
#   willow-serve worker-on [fast|batch|all]    # start worker lane(s)   (default: all)
#   willow-serve worker-off [fast|batch|all]   # stop worker lane(s)    (default: all)
#   willow-serve worker-status           # per-lane unit state
#   willow-serve worker-logs [fast|batch]      # follow a lane's logs   (default: fast)
#
# Port/host come from WILLOW_MCP_PORT / WILLOW_MCP_HOST (defaults 8766 / 127.0.0.1).
# The worker's app_id (whose confirmed 'tasks' mapping it drains) comes from
# WILLOW_APP_ID (default willow). Change them, re-run the install step, then on.
set -euo pipefail

UNIT="willow-mcp-serve"
PORT="${WILLOW_MCP_PORT:-8766}"
HOST="${WILLOW_MCP_HOST:-127.0.0.1}"
ENTRY_NAME="willow-mcp-serve"
WORKER_UNIT="willow-mcp-worker"
APP_ID="${WILLOW_APP_ID:-willow}"

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MCP_JSON="$REPO_ROOT/.mcp.json"
VENV_PY="$REPO_ROOT/.venv/bin/python3"
TEMPLATE="$REPO_ROOT/deploy/willow-mcp-serve.service.template"
WORKER_TEMPLATE="$REPO_ROOT/deploy/willow-mcp-worker@.service.template"
URL="http://$HOST:$PORT/mcp"
UNIT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/systemd/user"
UNIT_FILE="$UNIT_DIR/$UNIT.service"
WORKER_UNIT_FILE="$UNIT_DIR/$WORKER_UNIT@.service"

usage() {
  echo "usage: willow-serve {install|on|off|status|logs}" >&2
  echo "       willow-serve {worker-install|worker-on|worker-off|worker-status|worker-logs} [fast|batch|all]" >&2
  exit 2
}
have_unit() { [ -f "$UNIT_FILE" ]; }
have_worker_unit() { [ -f "$WORKER_UNIT_FILE" ]; }

# Expand a lane argument to instance unit names. all -> both lanes.
# Prints lanes to stdout; returns 2 on invalid input.
worker_lanes() {  # $1 = fast|batch|all (default all)
  case "${1:-all}" in
    fast|batch) echo "$1" ;;
    all) echo "fast batch" ;;
    *)
      echo "invalid lane: $1 (want fast|batch|all)" >&2
      return 2
      ;;
  esac
}

# Add or remove the http entry in .mcp.json. Idempotent; preserves other servers.
TOGGLE_PY="$REPO_ROOT/scripts/mcp_entry_toggle.py"
toggle_entry() {  # $1 = add|remove
  local py="python3"
  [ -x "$VENV_PY" ] && py="$VENV_PY"
  "$py" "$TOGGLE_PY" "$MCP_JSON" "$ENTRY_NAME" "$URL" "$1"
}

entry_present() { grep -q "\"$ENTRY_NAME\"" "$MCP_JSON" 2>/dev/null; }

case "${1:-}" in
  install)
    [ -f "$TEMPLATE" ] || { echo "missing template: $TEMPLATE" >&2; exit 1; }
    [ -x "$VENV_PY" ] || echo "warning: $VENV_PY not found — create the venv first (see README)" >&2
    mkdir -p "$UNIT_DIR"
    sed -e "s#@VENV@#$VENV_PY#g" \
        -e "s#@WORKDIR@#$REPO_ROOT#g" \
        -e "s#@PORT@#$PORT#g" \
        -e "s#@HOST@#$HOST#g" \
        "$TEMPLATE" > "$UNIT_FILE"
    systemctl --user daemon-reload
    echo "installed $UNIT_FILE  (port $PORT, host $HOST)"
    echo "next: willow-serve on"
    ;;
  on)
    have_unit || { echo "unit not installed — run: willow-serve install" >&2; exit 1; }
    systemctl --user start "$UNIT"
    toggle_entry add
    if systemctl --user is-active --quiet "$UNIT"; then
      echo "serve ON   $URL   (.mcp.json entry added — run /mcp to reconnect)"
    else
      echo "unit did not stay active — check: willow-serve logs" >&2
      exit 1
    fi
    ;;
  off)
    have_unit && systemctl --user stop "$UNIT" || true
    toggle_entry remove
    echo "serve OFF   (.mcp.json entry removed — run /mcp to reconnect)"
    ;;
  status)
    if have_unit && systemctl --user is-active --quiet "$UNIT"; then
      echo "unit:      active (port $PORT)"
    else
      echo "unit:      inactive"
    fi
    if entry_present; then
      echo ".mcp.json: entry present  ($URL)"
    else
      echo ".mcp.json: entry absent"
    fi
    ;;
  logs)
    journalctl --user -u "$UNIT" -f
    ;;
  worker-install)
    [ -f "$WORKER_TEMPLATE" ] || { echo "missing template: $WORKER_TEMPLATE" >&2; exit 1; }
    if [ -x "$VENV_PY" ]; then
      "$VENV_PY" -c "import kartikeya" 2>/dev/null \
        || echo "warning: kartikeya not importable from the venv — worker will exit at start (pip install -e . or pip install willow-mcp)" >&2
      if ! "$VENV_PY" -c "import mcp.server.mcpserver" 2>/dev/null; then
        echo "refreshing repo venv (MCP SDK 2.0 required for worker)…" >&2
        "$VENV_PY" -m pip install -e "$REPO_ROOT" -q
      fi
    else
      echo "warning: $VENV_PY not found — create the venv first (see README)" >&2
    fi
    # Fleet defaults — do not inherit the caller's WILLOW_HOME (hooks often set
    # the sovereign redirect). Override with WILLOW_SERVE_WILLOW_HOME if needed.
    WILLOW_HOME="${WILLOW_SERVE_WILLOW_HOME:-$HOME/github/.willow}"
    WILLOW_STORE_ROOT="${WILLOW_SERVE_STORE_ROOT:-$WILLOW_HOME/store}"
    WILLOW_PG_DB="${WILLOW_SERVE_PG_DB:-willow_20}"
    WILLOW_PG_USER="${WILLOW_SERVE_PG_USER:-${USER:-}}"
    HEARTBEAT_ROOT="${WILLOW_SERVE_HEARTBEAT_ROOT:-$WILLOW_HOME/worker_heartbeat}"
    KART_SANDBOX_CONFIG="${WILLOW_SERVE_KART_SANDBOX:-$WILLOW_HOME/kart-sandbox.json}"
    mkdir -p "$UNIT_DIR"
    sed -e "s#@VENV@#$VENV_PY#g" \
        -e "s#@WORKDIR@#$REPO_ROOT#g" \
        -e "s#@APPID@#$APP_ID#g" \
        -e "s#@WILLOW_HOME@#$WILLOW_HOME#g" \
        -e "s#@WILLOW_STORE_ROOT@#$WILLOW_STORE_ROOT#g" \
        -e "s#@WILLOW_PG_DB@#$WILLOW_PG_DB#g" \
        -e "s#@WILLOW_PG_USER@#$WILLOW_PG_USER#g" \
        -e "s#@HEARTBEAT_ROOT@#$HEARTBEAT_ROOT#g" \
        -e "s#@KART_SANDBOX_CONFIG@#$KART_SANDBOX_CONFIG#g" \
        "$WORKER_TEMPLATE" > "$WORKER_UNIT_FILE"
    systemctl --user daemon-reload
    echo "installed $WORKER_UNIT_FILE  (app_id $APP_ID; lanes are unit instances)"
    echo "  WILLOW_HOME=$WILLOW_HOME  PG=$WILLOW_PG_DB  sandbox=$KART_SANDBOX_CONFIG"
    echo "next: willow-serve worker-on   # starts fast + batch"
    ;;
  worker-on)
    have_worker_unit || { echo "unit not installed — run: willow-serve worker-install" >&2; exit 1; }
    lanes="$(worker_lanes "${2:-all}")" || exit $?
    for lane in $lanes; do
      systemctl --user start "$WORKER_UNIT@$lane"
      if systemctl --user is-active --quiet "$WORKER_UNIT@$lane"; then
        echo "worker ON   $lane lane"
      else
        echo "worker $lane did not stay active — check: willow-serve worker-logs $lane" >&2
        exit 1
      fi
    done
    ;;
  worker-off)
    lanes="$(worker_lanes "${2:-all}")" || exit $?
    for lane in $lanes; do
      have_worker_unit && systemctl --user stop "$WORKER_UNIT@$lane" || true
      echo "worker OFF  $lane lane"
    done
    ;;
  worker-status)
    for lane in fast batch; do
      if have_worker_unit && systemctl --user is-active --quiet "$WORKER_UNIT@$lane"; then
        echo "worker $lane:  active"
      else
        echo "worker $lane:  inactive"
      fi
    done
    ;;
  worker-logs)
    journalctl --user -u "$WORKER_UNIT@$(worker_lanes "${2:-fast}" | cut -d' ' -f1)" -f
    ;;
  *)
    usage
    ;;
esac
