#!/bin/bash
# Budget Pilot for AI coding tools (Codex CLI, Gemini CLI, Claude Code and others): install, update or remove.
#   curl -fsSL https://budgetpilot.net/install.sh | bash -s -- <licence key>     install or update (safe to run again)
#   curl -fsSL https://budgetpilot.net/install.sh | bash -s -- --remove          take it all out
# Needs: Linux or WSL (Intel/AMD or ARM) and curl. sudo and systemd are used when there; without them it installs for
# your user only (see below).
# What it does:
#   - downloads the pilot for this machine into ~/.local/share/budget-pilot/versions/<version>, with its bin pointing
#     at it (your key is checked before the download); later versions install themselves the same way (automatic
#     updates, signed by the seller; bp --autoupdate off stops them). Installs made before 3.5.1 stay in their
#     folder. The commands bp and budget-pilot are two-line starters in ~/.local/bin (on PATH for new terminals; bp
#     only when no other program has that name).
#   - activates your key: the seller's server gets the key, a machine code (a hash of /etc/machine-id), the host name
#     and the version - never anything from your chats - and answers with a signed lease for this machine
#   - Claude Code: puts the pilot's hooks + settings in Claude Code's machine-wide settings file
#     /etc/claude-code/managed-settings.json (other entries are kept; a backup copy goes to the pilot's folder). Without
#     sudo they go into your own ~/.claude/settings.json instead. Other CLIs: see `bp --detect`.
#   - runs the meter as the systemd service budget-pilot (starts at boot); without system systemd it uses a systemd
#     user service, or else starts it in the background and at login (crontab @reboot)
# Day to day:  bp (the Budget Pilot screen) | bp --status | --off | --on | --licence <key> | --help
# Multi-CLI settings: ~/.config/budget-pilot/config.json   Force a user-only install: BP_NO_SUDO=1
set -euo pipefail
SERVER="https://budgetpilot.net"
P_OLD=$HOME/.claude/pilot                                   # installs made before 3.5.1 stay there
if [ -n "${BUDGET_PILOT_DIR:-}" ]; then P=$BUDGET_PILOT_DIR
elif [ ! -e "${XDG_DATA_HOME:-$HOME/.local/share}/budget-pilot" ] && { [ -f "$P_OLD/licence.key" ] || [ -d "$P_OLD/versions" ]; }; then P=$P_OLD
else P=${XDG_DATA_HOME:-$HOME/.local/share}/budget-pilot; fi
MARK="# added by Budget Pilot"
B=$P/bin/budget-pilot
U=$(id -un)
G=$(id -gn)
UNIT=$HOME/.config/systemd/user/budget-pilot.service
STAMP=$(date +%Y%m%d-%H%M%S)
OS=linux; [ "$(uname -s)" = Darwin ] && OS=macos
AGENT=$HOME/Library/LaunchAgents/work.bpilot.budget-pilot.plist     # Mac: launchd keeps the meter running
[ "$OS" = macos ] && BP_NO_SUDO=1                           # Mac: Claude Code's own user settings, no password asked
WSL=""
grep -qi microsoft /proc/version 2>/dev/null && WSL=1
[ -n "${WSL_DISTRO_NAME:-}" ] && WSL=1

# how far can we go on this machine?
SUDO=""
# sudo -n first: cloud servers (Oracle, AWS, ...) give the user sudo without a password, but `sudo -v` still asks for one
if [ -z "${BP_NO_SUDO:-}" ] && command -v sudo >/dev/null && { sudo -n true 2>/dev/null || sudo -v 2>/dev/null; }; then SUDO=1; fi
SYSD=""
[ -n "$SUDO" ] && command -v systemctl >/dev/null && [ -d /run/systemd/system ] && SYSD=system
[ -z "$SYSD" ] && command -v systemctl >/dev/null && systemctl --user show-environment >/dev/null 2>&1 && SYSD=user

managed_setup() {
  sudo install -d -m 2775 -o root -g "$G" /etc/claude-code
  [ -f /etc/claude-code/managed-settings.json ] || echo '{}' | sudo tee /etc/claude-code/managed-settings.json >/dev/null
  sudo cp /etc/claude-code/managed-settings.json "$P/managed-settings.backup.$STAMP.json"
  sudo chown "$U" "$P/managed-settings.backup.$STAMP.json"
  sudo chgrp "$G" /etc/claude-code/managed-settings.json
  sudo chmod 664 /etc/claude-code/managed-settings.json
}

stop_meter() {
  [ -n "$SUDO" ] && sudo systemctl stop budget-pilot 2>/dev/null || true
  systemctl --user stop budget-pilot 2>/dev/null || true
  [ "$OS" = macos ] && { launchctl bootout "gui/$(id -u)" "$AGENT" 2>/dev/null || true; }
  pkill -f "$B daemon" 2>/dev/null || true
}

drop_cron() {
  if command -v crontab >/dev/null && crontab -l 2>/dev/null | grep -qF "$B daemon"; then
    crontab -l 2>/dev/null | grep -vF "$B daemon" | crontab - || true
  fi
}

if [ "${1:-}" = "--remove" ]; then
  mkdir -p "$P"
  stop_meter
  if [ -n "$SUDO" ]; then
    sudo systemctl disable budget-pilot 2>/dev/null || true
    [ -f /etc/claude-code/managed-settings.json ] && managed_setup
  fi
  [ -x "$B" ] && { "$B" off || true; }
  if [ -n "$SUDO" ] && [ -f /etc/systemd/system/budget-pilot.service ]; then
    sudo rm -f /etc/systemd/system/budget-pilot.service
    sudo systemctl daemon-reload || true
  fi
  if [ -f "$UNIT" ]; then
    systemctl --user disable budget-pilot 2>/dev/null || true
    rm -f "$UNIT"
    systemctl --user daemon-reload 2>/dev/null || true
  fi
  drop_cron
  rm -f "$AGENT"
  for S in "$HOME/.local/bin/budget-pilot" "$HOME/.local/bin/bp"; do
    if [ -f "$S" ] && grep -qF "$MARK" "$S"; then rm -f "$S"; fi
  done
  for f in "$HOME/.bashrc" "$HOME/.zshrc" "$HOME/.profile"; do
    if [ -f "$f" ] && grep -qF "$MARK" "$f"; then
      grep -vF "$MARK" "$f" > "$f.bp-tmp" || true; cat "$f.bp-tmp" > "$f"; rm -f "$f.bp-tmp"   # keeps links + rights
    fi
  done
  rm -rf "$P/bin" "$P/bin.old" "$P/versions"
  echo "Budget Pilot removed (its hooks and settings are out of your AI tools' settings; the meter is stopped)."
  echo "Your settings, log and memory (saved chat pieces) stay in $P and ~/.config/budget-pilot (delete them if you like)."
  exit 0
fi

command -v curl >/dev/null || { echo "Budget Pilot needs curl"; exit 1; }
KEY="${1:-}"
[ -z "$KEY" ] && [ -f "$P/licence.key" ] && KEY=$(cat "$P/licence.key")
[ -z "$KEY" ] && { echo "Usage: curl -fsSL $SERVER/install.sh | bash -s -- <licence key>  (get one at $SERVER)"; exit 1; }
case "$(uname -m)" in
  x86_64|amd64) ARCH=x86_64 ;;
  aarch64|arm64) ARCH=aarch64 ;;
  *) echo "Sorry, no Budget Pilot build for $(uname -m) machines yet."; exit 1 ;;
esac
command -v claude >/dev/null || command -v codex >/dev/null || command -v gemini >/dev/null || \
  echo "Note: no AI tool (Codex, Gemini, Claude Code) found yet; Budget Pilot finds it by itself once you install one."

K=""
[[ "$SERVER" =~ ^https://[0-9.]+(:[0-9]+)?$ ]] && K="-k"    # bare-IP server: self-signed certificate (the pilot checks signatures)
TMP=$(mktemp -d)
trap 'rm -rf "$TMP"' EXIT
DL=$ARCH; [ "$OS" = macos ] && DL=macos-$ARCH
echo "Downloading Budget Pilot ($DL)..."
CODE=$(curl $K -sS -o "$TMP/bp.tgz" -w '%{http_code}' --get --data-urlencode "key=$KEY" "$SERVER/download/$DL") || CODE=000
if [ "$CODE" != "200" ]; then
  echo "Download refused: $(cat "$TMP/bp.tgz" 2>/dev/null | head -c 300)"
  exit 1
fi
tar -xzf "$TMP/bp.tgz" -C "$TMP"
[ -x "$TMP/budget-pilot/budget-pilot" ] || { echo "The download looks broken; please try again."; exit 1; }
V=$(cat "$TMP/budget-pilot/VERSION" 2>/dev/null || "$TMP/budget-pilot/budget-pilot" version)
[[ "$V" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]] || { echo "The download looks broken; please try again."; exit 1; }

# layout: versions/<v>/ holds each version, bin -> versions/<v> (switched in one step, the hooks' path stays the same)
mkdir -p "$P/versions"
stop_meter
rm -rf "$P/versions/$V.new" && mv "$TMP/budget-pilot" "$P/versions/$V.new"
rm -rf "$P/versions/$V" && mv "$P/versions/$V.new" "$P/versions/$V"
rm -rf "$P/bin.old"
if [ -d "$P/bin" ] && [ ! -L "$P/bin" ]; then mv "$P/bin" "$P/bin.old"; fi    # 1.0.0 / 3.0.0: bin was a folder
if [ "$OS" = macos ]; then ln -sfn "versions/$V" "$P/bin"        # BSD mv has no -T; ln -n replaces the link itself
else ln -sfn "versions/$V" "$P/bin.new" && mv -Tf "$P/bin.new" "$P/bin"; fi
rm -rf "$P/bin.old"
[ -f "$P/bin/LICENSE.txt" ] && cp "$P/bin/LICENSE.txt" "$P/LICENSE.txt"

"$B" init-config "$SERVER" >/dev/null
"$B" licence "$KEY" || exit 1
[ -f "$HOME/.claude/settings.json" ] && cp "$HOME/.claude/settings.json" "$P/settings.backup.$STAMP.json"
if [ -n "$SUDO" ]; then
  # an earlier user-only install moves up to the machine-wide file (its user-file hooks are taken out by "on")
  "$B" config claude_settings >/dev/null
  if [ -f /etc/claude-code/managed-settings.json ] || command -v claude >/dev/null || [ -d "$HOME/.claude/projects" ]; then
    managed_setup
  fi
else
  "$B" config claude_settings user >/dev/null
  if command -v claude >/dev/null || [ -d "$HOME/.claude/projects" ]; then
    echo "No sudo: Claude Code's hooks go into your own Claude Code settings (this user only)."
  fi
fi
"$B" on
"$B" link --rc || true
CMD=budget-pilot     # 4.3.0 and newer also make the short "bp" (never over someone else's bp)
[ -x "$HOME/.local/bin/bp" ] && grep -qF "$MARK" "$HOME/.local/bin/bp" 2>/dev/null && CMD=bp
# Claude Code: say "on" only when its hooks really are in the settings file
if command -v claude >/dev/null || [ -d "$HOME/.claude/projects" ]; then
  if [ -n "$SUDO" ]; then CS=/etc/claude-code/managed-settings.json; else CS="$HOME/.claude/settings.json"; fi
  if grep -qF "pilot/bin/budget-pilot" "$CS" 2>/dev/null; then
    echo "Claude Code hooks are in $CS."
  else
    echo "Claude Code hooks could NOT be written to $CS. Fix it later with: $CMD --heal"
  fi
fi

if [ "$SYSD" = system ]; then
  sudo tee /etc/systemd/system/budget-pilot.service >/dev/null <<EOF
[Unit]
Description=Budget Pilot for AI coding CLIs: real usage meter, limit warnings, settings repair
After=network-online.target

[Service]
User=$U
Environment="HOME=$HOME"
ExecStart="$B" daemon
Restart=always
RestartSec=30
Nice=10

[Install]
WantedBy=multi-user.target
EOF
  sudo systemctl daemon-reload
  sudo systemctl enable budget-pilot >/dev/null 2>&1
  sudo systemctl restart budget-pilot
  drop_cron
elif [ "$SYSD" = user ]; then
  mkdir -p "$(dirname "$UNIT")"
  cat > "$UNIT" <<EOF
[Unit]
Description=Budget Pilot for AI coding CLIs: real usage meter, limit warnings, settings repair

[Service]
ExecStart="$B" daemon
Restart=always
RestartSec=30
Nice=10

[Install]
WantedBy=default.target
EOF
  systemctl --user daemon-reload
  systemctl --user enable budget-pilot >/dev/null 2>&1
  systemctl --user restart budget-pilot
  drop_cron
  echo "The meter runs as your user service (while you are logged in; 'loginctl enable-linger $U' keeps it on always)."
elif [ "$OS" = macos ]; then
  mkdir -p "$(dirname "$AGENT")"
  cat > "$AGENT" <<PLIST
<?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>work.bpilot.budget-pilot</string>
  <key>ProgramArguments</key><array><string>$B</string><string>daemon</string></array>
  <key>RunAtLoad</key><true/>
  <key>KeepAlive</key><true/>
  <key>ThrottleInterval</key><integer>30</integer>
  <key>ProcessType</key><string>Background</string>
</dict></plist>
PLIST
  launchctl bootout "gui/$(id -u)" "$AGENT" 2>/dev/null || true
  launchctl bootstrap "gui/$(id -u)" "$AGENT" 2>/dev/null || launchctl load -w "$AGENT" 2>/dev/null || true
  echo "The meter runs in the background and starts again when you log in."
else
  nohup "$B" daemon >/dev/null 2>&1 &
  if command -v crontab >/dev/null; then
    (crontab -l 2>/dev/null | grep -vF "$B daemon"; echo "@reboot $B daemon >/dev/null 2>&1") | crontab - 2>/dev/null || true
  fi
  if command -v crontab >/dev/null && crontab -l 2>/dev/null | grep -qF "@reboot $B daemon"; then
    echo "The meter runs in the background and starts again at boot (crontab)."
  else
    echo "The meter runs in the background until the next restart; start it again with: $B daemon &"
  fi
  if [ -n "$WSL" ]; then
    echo "On Windows (WSL) you can let it start with WSL: echo -e '[boot]\nsystemd=true' | sudo tee /etc/wsl.conf"
    echo "then, in PowerShell, wsl --shutdown, open Ubuntu again and rerun this command."
  fi
fi
sleep 3
IFS=. read -r VA VB _ <<< "$V"
if [ "$VA" -gt 4 ] || { [ "$VA" -eq 4 ] && [ "$VB" -ge 3 ]; }; then     # 4.3.0 and newer: the Budget Pilot screen
  "$B" screen
else
  "$B" status
  echo "All commands: budget-pilot help"
fi
if ! command -v budget-pilot >/dev/null 2>&1 && [ -x "$HOME/.local/bin/budget-pilot" ]; then
  echo "To use $CMD in THIS terminal, run this once (new terminals have it already):"
  echo '  export PATH="$HOME/.local/bin:$PATH"'
fi