#!/usr/bin/env bash
# omni-cswap installer — provision a machine (laptop or headless server) to run Omnigent
# with cswap multi-account load balancing.
#
# Idempotent: safe to re-run. Never touches credentials of accounts already added.
set -euo pipefail

OMNI_DIR="$HOME/.omni-cswap"
SERVER="https://agents.lepton.software"
# Where to fetch machinery/ when this script runs standalone (curl | bash). Defaults to the
# same host that serves the script; override with --base-url or OMNI_BASE_URL.
BASE_URL="${OMNI_BASE_URL:-https://setup.lepton.software}"
POOL=""; HEADLESS=0; BYPASS=0; SKIP_SERVICES=0; SKIP_HOST=0; CHECK_ONLY=0
LABEL_PREFIX="software.lepton"
MACHINERY_FILES="claude pick_account.py pretrust.py rescue.py report_usage.py"

# Piped to bash (`curl … | bash`), BASH_SOURCE is unset or "bash" and there is no script dir.
# Resolve to empty rather than guessing cwd, so machinery/ gets fetched instead of silently
# picking up an unrelated directory that happens to be named machinery.
SRC=""
if [ -n "${BASH_SOURCE[0]:-}" ] && [ -f "${BASH_SOURCE[0]:-}" ]; then
  SRC="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
fi

usage() {
  cat <<'EOF'
omni-cswap installer — provision a machine for Omnigent + cswap load balancing.

Remote (nothing to clone):
  curl -fsSL https://setup.lepton.software/install.sh | bash -s -- --headless --bypass
  bash <(curl -fsSL https://setup.lepton.software/install.sh)   # interactive (account logins)

Local:
  ./install.sh                    full install (interactive)

Flags:
  --headless         no browser logins; assumes accounts already added
  --pool "3 4 7"     set this machine's account allowlist
  --bypass           autonomous mode (headless hosts only)
  --server URL       Omnigent server (default https://agents.lepton.software)
  --base-url URL     where to fetch machinery/ from (env: OMNI_BASE_URL,
                     default https://setup.lepton.software)
  --skip-services    install files but don't register services
  --skip-host        register cswap services only; leave omnigent-host alone
                     (use when the host is supervised another way — a second
                      host double-registers with the server)
  --check            verify an existing install, change nothing
EOF
}

while [ $# -gt 0 ]; do
  case "$1" in
    --headless) HEADLESS=1; shift ;;
    --pool) POOL="$2"; shift 2 ;;
    --bypass) BYPASS=1; shift ;;
    --server) SERVER="$2"; shift 2 ;;
    --base-url) BASE_URL="${2%/}"; shift 2 ;;
    --skip-services) SKIP_SERVICES=1; shift ;;
    --skip-host) SKIP_HOST=1; shift ;;
    --check) CHECK_ONLY=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "unknown flag: $1" >&2; exit 2 ;;
  esac
done

# Under `curl … | bash`, stdin IS the script — a bare `read` would eat the remaining source
# instead of waiting for the user, and `claude /login` would get a pipe instead of a terminal.
# Bind fd 3 to the real terminal and read from it explicitly. Falls back to stdin when there
# is no tty (CI), where the interactive path is skipped anyway.
if [ -r /dev/tty ] && [ -t 1 ]; then exec 3</dev/tty; else exec 3<&0; fi

say()  { printf '\033[1;34m==>\033[0m %s\n' "$*"; }
ok()   { printf '  \033[32m✓\033[0m %s\n' "$*"; }
warn() { printf '  \033[33m!\033[0m %s\n' "$*"; }
die()  { printf '\033[31mERROR:\033[0m %s\n' "$*" >&2; exit 1; }

case "$(uname -s)" in
  Darwin) OS=mac ;;
  Linux)  OS=linux ;;
  *) die "unsupported OS: $(uname -s)" ;;
esac

# Tools live in ~/.local/bin (uv/native claude) or ~/.fnm/... (npm claude). A non-login shell —
# e.g. `ssh host './install.sh --check'` — sources no rc file, so those are absent from PATH and
# every check would report MISSING on a perfectly healthy box. Seed PATH before anything reads it.
export PATH="$HOME/.local/bin:$HOME/.fnm/aliases/default/bin:$PATH"

# ── 0. check mode ────────────────────────────────────────────────────────────
if [ "$CHECK_ONLY" = 1 ]; then
  say "Checking install"
  for b in uv claude cswap omni; do
    if command -v "$b" >/dev/null 2>&1; then ok "$b — $(command -v "$b")"; else warn "$b — MISSING"; fi
  done
  [ -x "$OMNI_DIR/claude" ] && ok "shim present" || warn "shim missing at $OMNI_DIR/claude"
  if [ -f "$OMNI_DIR/pool" ]; then ok "pool: $(tr '\n' ' ' <"$OMNI_DIR/pool")"; else warn "no pool file — this machine may use ANY account (see README)"; fi
  [ -f "$OMNI_DIR/bypass" ] && warn "bypass marker present — sessions run with bypassPermissions" || ok "bypass off (interactive approvals)"
  echo; say "Accounts"; cswap ls 2>&1 | sed 's/^/  /' || warn "cswap ls failed"
  echo; say "Shim resolution"
  ( export PATH="$OMNI_DIR:$PATH"; OMNI_CSWAP_DRYRUN=1 claude 2>&1 | sed 's/^/  /' ) || warn "shim dry-run failed"
  exit 0
fi

# ── 1. dependencies ──────────────────────────────────────────────────────────
say "Dependencies"

if ! command -v uv >/dev/null 2>&1; then
  curl -LsSf https://astral.sh/uv/install.sh | sh
  export PATH="$HOME/.local/bin:$PATH"
fi
command -v uv >/dev/null 2>&1 || die "uv install failed"
ok "uv — $(uv --version 2>/dev/null)"

# python3 is load-bearing: the shim runs pick_account.py under it, and the systemd/launchd units
# invoke rescue.py + report_usage.py as /usr/bin/python3. Usually pre-installed, but minimal
# images (bare containers, slim cloud rootfs) omit it — without it the load-balancing layer
# installs cleanly and then silently never routes a session. Ensure it, and that its sqlite3 +
# fcntl stdlib modules (report_usage / pick_account) actually import.
if ! command -v python3 >/dev/null 2>&1; then
  if [ "$OS" = linux ]; then
    sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y python3 || warn "install python3 manually"
  else
    warn "python3 missing — install Xcode Command Line Tools (xcode-select --install)"
  fi
fi
command -v python3 >/dev/null 2>&1 || die "python3 is required (machinery + services run on it)"
python3 -c "import sqlite3, fcntl, json" 2>/dev/null \
  && ok "python3 — $(python3 --version 2>&1 | awk '{print $2}') (sqlite3/fcntl ok)" \
  || die "python3 present but missing stdlib modules (need sqlite3, fcntl) — install a full python3"

if ! command -v tmux >/dev/null 2>&1; then
  # tmux is load-bearing: Omnigent runs each session in a tmux pane, and rescue.py drives them.
  if [ "$OS" = mac ]; then
    command -v brew >/dev/null 2>&1 && brew install tmux || warn "install tmux manually (no brew found)"
  else
    # DEBIAN_FRONTEND passed THROUGH sudo (sudo drops env by default): on a minimal image a dep
    # like tzdata would otherwise block on an interactive "Geographic area" prompt and hang.
    sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y tmux || warn "install tmux manually"
  fi
fi
command -v tmux >/dev/null 2>&1 && ok "tmux — $(tmux -V)" || warn "tmux MISSING — Omnigent sessions and rescue will not work"

# bubblewrap: Linux-only sandbox. Omnigent's own setup installs it and asks for confirmation,
# so we don't front-run it — we just report status.
if [ "$OS" = linux ]; then
  if command -v bwrap >/dev/null 2>&1; then ok "bubblewrap — $(bwrap --version)"
  else warn "bubblewrap not installed — Omnigent will prompt to install it on first run (confirm when asked)"; fi
else
  ok "bubblewrap — n/a on macOS"
fi

# ── 2. Claude Code CLI ───────────────────────────────────────────────────────
say "Claude Code CLI"
if ! command -v claude >/dev/null 2>&1; then
  curl -fsSL https://claude.ai/install.sh | bash
  export PATH="$HOME/.local/bin:$PATH"
fi
command -v claude >/dev/null 2>&1 || die "claude install failed — see https://claude.ai/install.sh"
ok "claude — $(claude --version 2>/dev/null | head -1)"

# ── 3. Omnigent + cswap ──────────────────────────────────────────────────────
say "Omnigent + cswap"
uv tool install --quiet omnigent 2>/dev/null || uv tool upgrade --quiet omnigent 2>/dev/null || true
uv tool install --quiet claude-swap 2>/dev/null || uv tool upgrade --quiet claude-swap 2>/dev/null || true
export PATH="$HOME/.local/bin:$PATH"
command -v omni  >/dev/null 2>&1 || die "omnigent install failed"
command -v cswap >/dev/null 2>&1 || die "claude-swap install failed"
ok "omnigent — $(command -v omni)"
ok "cswap — $(cswap --version 2>/dev/null)"

# ── 4. machinery ─────────────────────────────────────────────────────────────
say "Machinery -> $OMNI_DIR"

# Alongside the script (git clone / rsync)? Use it. Otherwise fetch from BASE_URL — that's the
# curl|bash path. Download to a temp dir FIRST and only copy into place once every file has
# arrived intact: a half-written shim on PATH breaks every session on the box.
if [ -n "$SRC" ] && [ -d "$SRC/machinery" ]; then
  MACHINERY="$SRC/machinery"
  ok "using local machinery ($MACHINERY)"
else
  command -v curl >/dev/null 2>&1 || die "curl is required to fetch machinery from $BASE_URL"
  MACHINERY="$(mktemp -d)"
  trap 'rm -rf "$MACHINERY"' EXIT
  say "Fetching machinery from $BASE_URL"
  for f in $MACHINERY_FILES; do
    curl -fsSL --retry 3 --retry-delay 1 "$BASE_URL/machinery/$f" -o "$MACHINERY/$f" \
      || die "could not fetch $BASE_URL/machinery/$f (set --base-url ?)"
    [ -s "$MACHINERY/$f" ] || die "fetched $f is empty — refusing to install a truncated file"
  done
  # Cheap integrity gate: every file must be the kind of file we expect. Catches a captive-portal
  # HTML page or an error body served with 200, which curl -f would happily accept.
  head -1 "$MACHINERY/claude" | grep -q '^#!' || die "fetched shim is not a script — bad BASE_URL?"
  for f in pick_account.py pretrust.py rescue.py report_usage.py; do
    head -1 "$MACHINERY/$f" | grep -q '^#!.*python' || die "fetched $f is not a python script"
  done
  ok "fetched $(echo $MACHINERY_FILES | wc -w | tr -d ' ') files"
fi

mkdir -p "$OMNI_DIR"
for f in $MACHINERY_FILES; do
  cp "$MACHINERY/$f" "$OMNI_DIR/$f"
done
chmod +x "$OMNI_DIR/claude" "$OMNI_DIR"/*.py
ok "shim + picker + pretrust + rescue + reporter installed"

if [ -n "$POOL" ]; then
  printf '%s\n' $POOL > "$OMNI_DIR/pool"
  ok "pool = $POOL"
elif [ ! -f "$OMNI_DIR/pool" ]; then
  warn "no pool file — this machine may pick ANY account."
  warn "  Two machines sharing an account WILL kill each other's OAuth refresh token."
  warn "  Set a disjoint pool per machine: ./install.sh --pool \"3 4\""
fi

if [ "$BYPASS" = 1 ]; then
  touch "$OMNI_DIR/bypass"
  warn "bypass marker set — sessions on this host run --permission-mode bypassPermissions"
else
  [ -f "$OMNI_DIR/bypass" ] && warn "bypass marker already present (remove $OMNI_DIR/bypass to disable)"
fi

# PATH: the shim dir must come FIRST so Omnigent's `claude` resolves to it.
for rc in "$HOME/.zshrc" "$HOME/.bashrc"; do
  [ -f "$rc" ] || continue
  grep -q '.omni-cswap' "$rc" 2>/dev/null && continue
  printf '\n# omni-cswap: shim must precede real claude on PATH\nexport PATH="$HOME/.omni-cswap:$HOME/.local/bin:$PATH"\n' >> "$rc"
  ok "PATH updated in $(basename "$rc")"
done

# ── 5. accounts ──────────────────────────────────────────────────────────────
say "Accounts"
# INVARIANT: OAuth refresh tokens are single-use. Log in on THIS machine, per account.
# Never `cswap export` on one box and `import` on another — both copies then race to refresh
# the same grant and mutually log each other out.
if [ "$HEADLESS" = 1 ]; then
  warn "--headless: skipping interactive logins"
  cswap ls 2>&1 | sed 's/^/  /' || true
else
  echo "  Log in each account on THIS machine (never export/import between machines)."
  echo "  For each: a browser opens -> log in -> then it's registered as a cswap slot."
  while true; do
    cswap ls 2>&1 | sed 's/^/  /' || true
    echo
    read -rp "  Add an account on this machine? [y/N] " yn <&3
    case "$yn" in
      [Yy]*) ;;
      *) break ;;
    esac
    read -rp "  Slot number (blank = next free): " slot <&3
    say "Opening Claude Code login — complete it in the browser, then exit claude (Ctrl-D)."
    # Bypass the shim for auth: utility invocations must hit the real claude/default config.
    # stdin comes from fd 3 (the real terminal) so this works under `curl … | bash` too.
    ( PATH="$(printf '%s' "$PATH" | tr ':' '\n' | grep -vxF "$OMNI_DIR" | paste -sd: -)"; claude /login <&3 || claude <&3 )
    if [ -n "$slot" ]; then cswap add --slot "$slot" || warn "cswap add failed"
    else cswap add || warn "cswap add failed"; fi
  done
fi

# ── 6. auto-rotation settings ────────────────────────────────────────────────
say "Auto-rotation"
# Rotate before the active account hits its window, and count Fable's per-model weekly limit
# (Omnigent leans on Fable; without this, a Fable-capped account still looks "healthy" at the
# account level and auto never rotates off it).
cswap config set autoswitch.threshold 90 >/dev/null 2>&1 || true
cswap config set autoswitch.strategy best >/dev/null 2>&1 || true
cswap config set autoswitch.model Fable >/dev/null 2>&1 || true
ok "threshold=90 strategy=best model=Fable"
warn "cswap auto reads settings at start — restart the service after changing them"

# ── 7. services ──────────────────────────────────────────────────────────────
if [ "$SKIP_SERVICES" = 1 ]; then
  warn "--skip-services: not registering background services"
else
say "Services"
# Shim dir FIRST so Omnigent's `claude` resolves to it. The fnm path matters on boxes where
# Claude Code came from npm (`@anthropic-ai/claude-code`) rather than the native installer —
# without it the shim's `exec claude` finds nothing. Harmless when absent.
SHIM_PATH="$OMNI_DIR:$HOME/.fnm/aliases/default/bin:$HOME/.local/bin:/opt/homebrew/bin:/opt/homebrew/sbin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"
mkdir -p "$HOME/.omnigent/logs"

if [ "$OS" = mac ]; then
  LA="$HOME/Library/LaunchAgents"; mkdir -p "$LA"
  emit_plist() { # label, program-args-xml
    cat > "$LA/$1.plist" <<EOF
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0"><dict>
 <key>Label</key><string>$1</string>
 <key>ProgramArguments</key><array>$2</array>
 <key>EnvironmentVariables</key><dict><key>PATH</key><string>$SHIM_PATH</string></dict>
 <key>RunAtLoad</key><true/><key>KeepAlive</key><true/><key>ThrottleInterval</key><integer>30</integer>
 <key>StandardOutPath</key><string>$3</string>
 <key>StandardErrorPath</key><string>$3</string>
</dict></plist>
EOF
  }
  emit_plist "$LABEL_PREFIX.cswap-auto" \
    "<string>$HOME/.local/bin/cswap</string><string>auto</string>" \
    "$OMNI_DIR/cswap-auto.log"
  emit_plist "$LABEL_PREFIX.cswap-rescue" \
    "<string>/usr/bin/python3</string><string>$OMNI_DIR/rescue.py</string>" \
    "$OMNI_DIR/rescue.log"
  emit_plist "$LABEL_PREFIX.cswap-report" \
    "<string>/usr/bin/python3</string><string>$OMNI_DIR/report_usage.py</string>" \
    "$OMNI_DIR/report.log"
  SERVICES="cswap-auto cswap-rescue cswap-report"
  if [ "$SKIP_HOST" = 1 ]; then
    warn "--skip-host: leaving omnigent-host to its existing supervision"
  else
    emit_plist "$LABEL_PREFIX.omnigent-host" \
      "<string>$HOME/.local/bin/omni</string><string>host</string><string>--server</string><string>$SERVER</string><string>--non-interactive</string>" \
      "$HOME/.omnigent/logs/launchagent-host.out.log"
    SERVICES="$SERVICES omnigent-host"
  fi

  for svc in $SERVICES; do
    L="$LABEL_PREFIX.$svc"
    # bootout-then-bootstrap, never kickstart: kickstart reuses the OLD plist, so config
    # changes silently don't take effect.
    launchctl bootout "gui/$UID/$L" 2>/dev/null || true
    launchctl bootstrap "gui/$UID" "$LA/$L.plist" 2>/dev/null || warn "bootstrap $L failed"
    ok "$L"
  done
else
  SD="$HOME/.config/systemd/user"; mkdir -p "$SD"
  emit_unit() { # name, description, exec, log
    cat > "$SD/$1.service" <<EOF
[Unit]
Description=$2
After=network-online.target

[Service]
Type=simple
Environment=PATH=$SHIM_PATH
ExecStart=$3
Restart=always
RestartSec=30
StandardOutput=append:$4
StandardError=append:$4

[Install]
WantedBy=default.target
EOF
  }
  emit_unit cswap-auto   "cswap auto-rotation"    "$HOME/.local/bin/cswap auto" "$OMNI_DIR/cswap-auto.log"
  emit_unit cswap-rescue "cswap rescue watchdog"  "/usr/bin/python3 $OMNI_DIR/rescue.py" "$OMNI_DIR/rescue.log"
  emit_unit cswap-report "cswap usage reporter"   "/usr/bin/python3 $OMNI_DIR/report_usage.py" "$OMNI_DIR/report.log"
  SERVICES="cswap-auto cswap-rescue cswap-report"
  if [ "$SKIP_HOST" = 1 ]; then
    warn "--skip-host: leaving omnigent-host to its existing supervision"
  else
    emit_unit omnigent-host "Omnigent host runner" "$HOME/.local/bin/omni host --server $SERVER --non-interactive" "$HOME/.omnigent/logs/host.log"
    SERVICES="$SERVICES omnigent-host"
  fi

  systemctl --user daemon-reload
  for svc in $SERVICES; do
    systemctl --user enable --now "$svc.service" 2>/dev/null && ok "$svc" || warn "$svc failed to start"
  done
  # Headless boxes: without lingering, --user services die when your SSH session ends.
  loginctl enable-linger "$USER" 2>/dev/null && ok "linger enabled (services survive logout)" \
    || warn "could not enable linger — run: sudo loginctl enable-linger $USER"
fi
fi

# ── 8. verify ────────────────────────────────────────────────────────────────
echo; say "Verify"
( export PATH="$OMNI_DIR:$PATH"; OMNI_CSWAP_DRYRUN=1 claude 2>&1 | sed 's/^/  /' ) \
  || warn "shim dry-run failed — check $OMNI_DIR/shim.log"
OMNI_REPORT_ONCE=1 python3 "$OMNI_DIR/report_usage.py" 2>&1 | sed 's/^/  /' \
  || warn "usage reporter failed"

echo
say "Done."
echo "  Accounts:   cswap ls"
if [ -n "$SRC" ]; then
  echo "  Health:     $SRC/install.sh --check"
else
  echo "  Health:     curl -fsSL $BASE_URL/install.sh | bash -s -- --check"
fi
echo "  Usage data: $OMNI_DIR/usage/usage.db  (sqlite3 … 'select * from usage_snapshots')"
echo "  Logs:       $OMNI_DIR/{shim,rescue,cswap-auto,report}.log"
[ -f "$OMNI_DIR/pool" ] || echo "  NOTE: set a per-machine pool before running a second machine (see README)."