#!/usr/bin/env bash
#
# install-client.sh — Provision (or update) a pi client machine with the
# parts of pi-extensions you select. Idempotent, non-destructive by default
# only when --update-only is set.
#
# Run from a checked-out copy of the pi-extensions repo:
#
#   scripts/install-client.sh                          # everything
#   scripts/install-client.sh --components ai-server,certs
#   scripts/install-client.sh --update-only            # add new files only,
#                                                      # never replace existing
#
# Components:
#   certs            mTLS client cert + root CA from the Caddy host
#   ssh              SSH key-auth to the ai-server host (admin commands)
#   ai-server        the pi extension (mTLS llama.cpp router provider)
#   dark-mechanicus  theme TUI bundle (banner, indicator, status line, etc.)
#   token-stats      footer owner for context + token rate display
#   memory           cross-agent memory access extension (memory/)
#   themes           JSON theme palette files (themes/*.json)
#   local-llama      split local llama.cpp extension (llama.cpp/)
#   ai-complete      shell CLI for direct llama-server access (installed to
#                    --ai-complete-dir, default ~/bin/)
#   shared           shared TS utils — auto-included if dark-mechanicus is
#   all              shorthand for every component above (default if neither
#                    --components nor any --no-<x> flag is given)
#
# Modes:
#   --update-only    Do not replace existing files. Only add new ones. Useful
#                    after a `git pull` to surface new modules without
#                    clobbering local edits.
#   --force          Overwrite existing cert files (other components: same as
#                    no flag — rsync replaces unless --update-only).
#
# Options:
#   --components LIST           Comma-separated component names. Default: all.
#   --update-only               Non-destructive: only add missing files.
#   --ai-complete-dir PATH      Where to install ai-complete (default: ~/bin)
#   --caddy-host USER@HOST      Default: shahondin1624@192.168.2.2
#   --caddy-cert-dir PATH       Default: /mnt/ssdpool/@docker/caddy/certs
#   --ai-server-host USER@HOST  Default: ai-server@192.168.2.3
#   --ai-server-url URL         Default: https://ai.shahondin1624.de
#   --pi-dir PATH               Default: $HOME/.pi/agent
#   --no-certs                  Same as omitting "certs" from --components
#   --no-ssh-setup              Same as omitting "ssh" from --components
#   --force                     Overwrite existing cert files
#   --skip-verify               Skip the final mTLS health probe
#   --help                      Show this message

set -euo pipefail

# ─── defaults ────────────────────────────────────────────────────────────

CADDY_HOST="${CADDY_HOST:-shahondin1624@192.168.2.2}"
CADDY_CERT_DIR="${CADDY_CERT_DIR:-/mnt/ssdpool/@docker/caddy/certs}"
AI_SERVER_HOST="${AI_SERVER_HOST:-ai-server@192.168.2.3}"
AI_SERVER_URL="${AI_SERVER_URL:-https://ai.shahondin1624.de}"
PI_DIR="${PI_DIR:-$HOME/.pi/agent}"
AI_COMPLETE_DIR="${AI_COMPLETE_DIR:-$HOME/bin}"
COMPONENTS=""
UPDATE_ONLY=0
FORCE=0
SKIP_VERIFY=0
LEGACY_NO_CERTS=0
LEGACY_NO_SSH=0

ALL_COMPONENTS=(certs ssh ai-server dark-mechanicus token-stats memory themes local-llama ai-complete shared)

usage() { sed -n '2,/^$/p' "$0" | sed 's/^#\{0,1\} \{0,1\}//'; exit 0; }

while [[ $# -gt 0 ]]; do
	case "$1" in
		--components)        COMPONENTS="$2"; shift 2 ;;
		--update-only)       UPDATE_ONLY=1; shift ;;
		--ai-complete-dir)   AI_COMPLETE_DIR="$2"; shift 2 ;;
		--caddy-host)        CADDY_HOST="$2"; shift 2 ;;
		--caddy-cert-dir)    CADDY_CERT_DIR="$2"; shift 2 ;;
		--ai-server-host)    AI_SERVER_HOST="$2"; shift 2 ;;
		--ai-server-url)     AI_SERVER_URL="$2"; shift 2 ;;
		--pi-dir)            PI_DIR="$2"; shift 2 ;;
		--no-ssh-setup)      LEGACY_NO_SSH=1; shift ;;
		--no-certs)          LEGACY_NO_CERTS=1; shift ;;
		--force)             FORCE=1; shift ;;
		--skip-verify)       SKIP_VERIFY=1; shift ;;
		-h|--help)           usage ;;
		*) echo "Unknown arg: $1 (try --help)" >&2; exit 1 ;;
	esac
done

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$REPO_ROOT"
[[ -d ai-server ]] || { echo "Run from a checked-out pi-extensions repo (ai-server/ not found here)" >&2; exit 1; }

# ─── component selection ─────────────────────────────────────────────────

# If --components not given, default to all, then strip what --no-* removed
if [[ -z "$COMPONENTS" ]]; then
	COMPONENTS="all"
fi

# Expand "all" -> every known component
if [[ "$COMPONENTS" == "all" || "$COMPONENTS" == *",all,"* || "$COMPONENTS" == "all,"* || "$COMPONENTS" == *",all" ]]; then
	COMPONENTS=$(IFS=,; echo "${ALL_COMPONENTS[*]}")
fi

# Honor legacy --no-certs / --no-ssh-setup by removing from the list
if [[ $LEGACY_NO_CERTS -eq 1 ]]; then
	COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | grep -vx 'certs' | paste -sd, -)
fi
if [[ $LEGACY_NO_SSH -eq 1 ]]; then
	COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | grep -vx 'ssh' | paste -sd, -)
fi

# Auto-add 'shared' if dark-mechanicus, token-stats, or memory is selected (they import from ../shared/ or ./shared/)
if [[ ",$COMPONENTS," == *",dark-mechanicus,"* || ",$COMPONENTS," == *",token-stats,"* || ",$COMPONENTS," == *",memory,"* ]]; then
	if [[ ",$COMPONENTS," != *",shared,"* ]]; then
	COMPONENTS="$COMPONENTS,shared"
	fi
fi

# Normalize and dedupe
COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | sed '/^$/d' | awk '!seen[$0]++' | paste -sd, -)

want() { [[ ",$COMPONENTS," == *",$1,"* ]]; }

# Validate component names
for c in $(echo "$COMPONENTS" | tr ',' ' '); do
	found=0
	for known in "${ALL_COMPONENTS[@]}"; do
		[[ "$c" == "$known" ]] && { found=1; break; }
	done
	[[ $found -eq 1 ]] || { echo "Unknown component: '$c' (valid: ${ALL_COMPONENTS[*]}, all)" >&2; exit 1; }
done

# rsync flags differ per mode
if [[ $UPDATE_ONLY -eq 1 ]]; then
	RSYNC_FLAGS=(-a --ignore-existing)
	MODE_LABEL="update-only (existing files preserved)"
else
	RSYNC_FLAGS=(-a)
	MODE_LABEL="full sync (existing files replaced)"
fi

echo "==> pi-extensions client install"
echo "    repo root:    $REPO_ROOT"
echo "    pi dir:       $PI_DIR"
echo "    components:   $COMPONENTS"
echo "    mode:         $MODE_LABEL"
[[ $FORCE -eq 1 ]] && echo "    cert force:   on (will overwrite existing certs)"

mkdir -p "$PI_DIR/extensions"

# ─── components ──────────────────────────────────────────────────────────

if want certs; then
	echo; echo "==> [certs] installing mTLS certs from $CADDY_HOST"
	mkdir -p "$PI_DIR/certs"
	chmod 700 "$PI_DIR/certs"
	for pair in "client.crt:client.pem" "client.key:client-key.pem" "root-ca.pem:root-ca.pem"; do
		src="${pair%%:*}"; dst="${pair##*:}"
		target="$PI_DIR/certs/$dst"
		# Cert preservation rules:
		#   - default                  : keep existing (don't clobber a working identity)
		#   - --force                  : overwrite
		#   - --update-only            : keep existing (same as default)
		if [[ -f "$target" && $FORCE -eq 0 ]]; then
			echo "    keep existing $dst (use --force to replace)"
			continue
		fi
		echo "    fetch $CADDY_HOST:$CADDY_CERT_DIR/$src -> $target"
		scp -q "$CADDY_HOST:$CADDY_CERT_DIR/$src" "$target"
	done
	chmod 600 "$PI_DIR"/certs/*.pem 2>/dev/null || true
fi

if want ai-server; then
	echo; echo "==> [ai-server] syncing extension"
	rsync "${RSYNC_FLAGS[@]}" --exclude='.git' ai-server "$PI_DIR/extensions/"
fi

if want dark-mechanicus; then
	if [[ -d dark-mechanicus ]]; then
		echo; echo "==> [dark-mechanicus] syncing theme TUI bundle"
		rsync "${RSYNC_FLAGS[@]}" --exclude='.git' dark-mechanicus "$PI_DIR/extensions/"
	else
		echo "    (dark-mechanicus/ not in repo, skipping)"
	fi
fi

if want token-stats; then
	if [[ -d token-stats ]]; then
		echo; echo "==> [token-stats] syncing footer extension"
		rsync "${RSYNC_FLAGS[@]}" --exclude='.git' token-stats "$PI_DIR/extensions/"
	else
		echo "    (token-stats/ not in repo, skipping)"
	fi
fi

if want memory; then
	if [[ -d memory ]]; then
		echo; echo "==> [memory] syncing cross-agent memory extension"
		rsync "${RSYNC_FLAGS[@]}" --exclude='.git' memory "$PI_DIR/extensions/"
	else
		echo "    (memory/ not in repo, skipping)"
	fi
fi

if want shared; then
	if [[ -d shared ]]; then
		echo; echo "==> [shared] syncing shared TS utils"
		rsync "${RSYNC_FLAGS[@]}" --exclude='.git' shared "$PI_DIR/extensions/"
	else
		echo "    (shared/ not in repo, skipping)"
	fi
fi

if want themes; then
	if [[ -d themes ]]; then
		echo; echo "==> [themes] syncing palette files"
		mkdir -p "$PI_DIR/themes"
		rsync "${RSYNC_FLAGS[@]}" themes/ "$PI_DIR/themes/"
		echo "    $(ls themes/*.json 2>/dev/null | wc -l) theme file(s) installed"
	fi
fi

if want local-llama; then
	if [[ -d llama.cpp ]]; then
		echo; echo "==> [local-llama] syncing llama.cpp extension"
		rsync "${RSYNC_FLAGS[@]}" --exclude='.git' llama.cpp "$PI_DIR/extensions/"
	else
		echo "    (llama.cpp/ not in repo, skipping)"
	fi
fi

if want ai-complete; then
	echo; echo "==> [ai-complete] installing CLI to $AI_COMPLETE_DIR"
	mkdir -p "$AI_COMPLETE_DIR"
	target="$AI_COMPLETE_DIR/ai-complete"
	if [[ -f "$target" && $UPDATE_ONLY -eq 1 ]]; then
		echo "    keep existing $target (--update-only)"
	else
		install -m 755 scripts/ai-complete "$target"
		echo "    installed $target"
	fi
	# Add to PATH in shell rc if missing (idempotent)
	case ":$PATH:" in
		*":$AI_COMPLETE_DIR:"*) ;;
		*)
			rc=""
			[[ -f "$HOME/.bashrc" ]] && rc="$HOME/.bashrc"
			[[ -f "$HOME/.zshrc"  ]] && rc="$HOME/.zshrc"
			if [[ -n "$rc" ]] && ! grep -qF "$AI_COMPLETE_DIR" "$rc" 2>/dev/null; then
				echo "    appending PATH += $AI_COMPLETE_DIR to $rc"
				echo "export PATH=\"$AI_COMPLETE_DIR:\$PATH\"" >> "$rc"
			fi
			;;
	esac
fi

if want ssh; then
	echo; echo "==> [ssh] checking SSH key-auth to $AI_SERVER_HOST"
	if ssh -o BatchMode=yes -o ConnectTimeout=5 "$AI_SERVER_HOST" 'true' 2>/dev/null; then
		echo "    key-auth already works"
	else
		if [[ ! -f "$HOME/.ssh/id_ed25519.pub" && ! -f "$HOME/.ssh/id_rsa.pub" ]]; then
			echo "    no local ssh key; generating ed25519"
			ssh-keygen -t ed25519 -N '' -f "$HOME/.ssh/id_ed25519"
		fi
		echo "    running ssh-copy-id (will prompt for password once)"
		ssh-copy-id "$AI_SERVER_HOST"
	fi
fi

# ─── verification ────────────────────────────────────────────────────────

if [[ $SKIP_VERIFY -eq 0 ]] && want certs; then
	echo; echo "==> Verifying mTLS reachability"
	# Server cert is LE-issued (publicly trusted) — rely on the system CA
	# bundle, matching ai-server/{stream,admin}.ts which omit `ca:`. The local
	# root-ca.pem is kept around only to satisfy loadCerts() in config.ts.
	if curl -sS --max-time 5 \
			--cert "$PI_DIR/certs/client.pem" \
			--key "$PI_DIR/certs/client-key.pem" \
			"$AI_SERVER_URL/health" | grep -q '"ok"'; then
		echo "    ✓ mTLS handshake + /health OK"
	else
		echo "    ✗ could not verify — check cert files and server reachability"
		exit 1
	fi
fi

echo
echo "==> Done."
echo "   Next: open pi and /reload, or run 'pi' fresh."
want ai-complete && echo "   ai-complete installed; open a new shell (or 'hash -r') to pick up PATH."