#!/usr/bin/env bash
# install.sh — stand up the whole vm_mgr stack (the station-console + on-demand
# VM creation) on a fresh host. Idempotent: safe to re-run; each step checks
# before it acts. Run as root:  sudo ./install.sh   (reads ./install.conf)
#
# What it does (each step a no-op if already satisfied):
#   1. LXD (snap) + minimal `lxd init`
#   2. btrfs storage pool on POOL_DISK (sized, not over-committed)
#   3. profiles: `default` (NAT via lxdbr0) and `isolated` (no shares/nic)
#   4. golden image: import GOLDEN_TARBALL or build from BASE_IMAGE, bake llm-station
#   5. vm_mgr bin/ on PATH (symlinks into /usr/local/bin)
#   6. console: python venv + aiohttp, self-signed TLS, admin password, systemd unit
#   7. enable + start station-console
#
# Flags: --yes (don't prompt), --skip-golden, --skip-console, --dry-run, -h
set -euo pipefail

ROOT="$(cd "$(dirname "$(readlink -f "$0")")" && pwd)"
CONF="$ROOT/install.conf"
ASSUME_YES=0; SKIP_GOLDEN=0; SKIP_CONSOLE=0; DRY=0
while [ "$#" -gt 0 ]; do
  case "$1" in
    -y|--yes) ASSUME_YES=1;;
    --skip-golden) SKIP_GOLDEN=1;;
    --skip-console) SKIP_CONSOLE=1;;
    --dry-run) DRY=1;;
    -h|--help) grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0;;
    *) echo "unknown option: $1" >&2; exit 2;;
  esac
  shift
done

# ---- helpers ----------------------------------------------------------------
c_grn=$'\e[32m'; c_yel=$'\e[33m'; c_red=$'\e[31m'; c_rst=$'\e[0m'
log()  { echo "${c_grn}==>${c_rst} $*"; }
warn() { echo "${c_yel}warn:${c_rst} $*" >&2; }
die()  { echo "${c_red}error:${c_rst} $*" >&2; exit 1; }
run()  { if [ "$DRY" = 1 ]; then echo "  [dry-run] $*"; else "$@"; fi; }
have() { command -v "$1" >/dev/null 2>&1; }
confirm() {
  [ "$ASSUME_YES" = 1 ] && return 0
  read -r -p "$1 [y/N] " a; [[ "$a" =~ ^[Yy]$ ]]
}
# Render the shipped systemd unit templates (*.service.in) into concrete
# *.service files, substituting the install root / run-as user / python path.
# It does NOT enable anything — you cp the ones you want into /etc/systemd/system.
render_units() {
  local py="$ROOT/console/.venv/bin/python3"
  [ -x "$py" ] || py="/usr/bin/python3"
  local t out n=0
  while IFS= read -r -d '' t; do
    out="${t%.in}"
    if [ "$DRY" = 1 ]; then
      echo "  [dry-run] would render $(basename "$out")"
    else
      sed -e "s|@VM_MGR_ROOT@|$ROOT|g" \
          -e "s|@CONSOLE_USER@|${CONSOLE_USER:-root}|g" \
          -e "s|@PYTHON@|$py|g" "$t" > "$out"
    fi
    n=$((n + 1))
  done < <(find "$ROOT" -name '*.service.in' -print0)
  log "Rendered $n systemd unit template(s): *.service.in -> *.service (not enabled; cp the ones you want into /etc/systemd/system/)"
}

# ---- preflight --------------------------------------------------------------
[ "$(id -u)" = 0 ] || die "run as root (sudo ./install.sh)"
have systemctl || die "systemd required"
[ -f "$CONF" ] || die "missing $CONF — copy/edit the template first"
# shellcheck source=/dev/null
. "$CONF"

# defaults for anything left blank
POOL_NAME="${POOL_NAME:-sandbox}"
POOL_SIZE="${POOL_SIZE:-400GiB}"
GOLDEN_NAME="${GOLDEN_NAME:-sandbox}"
BASE_IMAGE="${BASE_IMAGE:-ubuntu:24.04}"
CONSOLE_HOST="${CONSOLE_HOST:-0.0.0.0}"
CONSOLE_PORT="${CONSOLE_PORT:-8800}"
CONSOLE_DOMAIN="${CONSOLE_DOMAIN:-console.local}"
STATION_UID="${STATION_UID:-1000}"; STATION_GID="${STATION_GID:-1000}"
STATION_HOME="${STATION_HOME:-/home/ubuntu}"
# run-as user: explicit, else the invoking sudo user, else fail (never root)
CONSOLE_USER="${CONSOLE_USER:-${SUDO_USER:-}}"
[ -n "$CONSOLE_USER" ] || die "set CONSOLE_USER in install.conf (won't run the console as root)"
id "$CONSOLE_USER" >/dev/null 2>&1 || die "CONSOLE_USER '$CONSOLE_USER' does not exist"
[ -n "${POOL_DISK:-}" ] || die "set POOL_DISK in install.conf"

log "vm_mgr installer  (root=$ROOT, pool=$POOL_NAME on $POOL_DISK, console as $CONSOLE_USER:lxd:$CONSOLE_PORT)"

# ---- 1. LXD -----------------------------------------------------------------
if ! have lxc; then
  log "Installing LXD (snap)"
  have snap || run apt-get update -qq && run apt-get install -y snapd
  run snap install lxd
  export PATH="/snap/bin:$PATH"
fi
have lxc || export PATH="/snap/bin:$PATH"
if ! lxc storage list >/dev/null 2>&1; then
  log "Initialising LXD (lxd init --auto)"
  run lxd init --auto
fi
# the console's run-as user must be able to reach the LXD socket
if ! id -nG "$CONSOLE_USER" | tr ' ' '\n' | grep -qx lxd; then
  log "Adding $CONSOLE_USER to the lxd group"
  run usermod -aG lxd "$CONSOLE_USER"
  warn "$CONSOLE_USER gained lxd group — its login sessions must re-auth to see it (the service picks it up on start)"
fi

# ---- 2. storage pool --------------------------------------------------------
if lxc storage show "$POOL_NAME" >/dev/null 2>&1; then
  log "Storage pool '$POOL_NAME' exists"
else
  log "Creating btrfs pool '$POOL_NAME' ($POOL_SIZE) under $POOL_DISK"
  run mkdir -p "$POOL_DISK"
  run lxc storage create "$POOL_NAME" btrfs source="$POOL_DISK/${POOL_NAME}.img" size="$POOL_SIZE"
fi

# ---- 3. profiles ------------------------------------------------------------
ensure_profile() {  # name, shared(0/1) — shared adds NAT nic + RO shared-plane mounts
  local p="$1" shared="$2"
  lxc profile show "$p" >/dev/null 2>&1 || run lxc profile create "$p"
  run lxc profile device add "$p" root disk path=/ pool="$POOL_NAME" 2>/dev/null || true
  if [ "$shared" = 1 ]; then
    run lxc profile device add "$p" eth0 nic network=lxdbr0 name=eth0 2>/dev/null || true
    # RO shared plane (host-editable, read-only in the VM). The 'scripts' mount
    # also carries share/scripts/directives — the home for linked keeper directives.
    for s in modules scripts services; do
      run mkdir -p "$ROOT/share/$s"
      run lxc profile device add "$p" "common-$s" disk \
        source="$ROOT/share/$s" path="/srv/share/$s" readonly=true 2>/dev/null || true
    done
  fi
}
log "Ensuring profiles: default (NAT + shared plane), isolated (no nic, no shares)"
ensure_profile default 1
ensure_profile isolated 0

# ---- 4. golden image --------------------------------------------------------
if [ "$SKIP_GOLDEN" = 1 ]; then
  log "Skipping golden image (--skip-golden)"
elif lxc info "$GOLDEN_NAME" >/dev/null 2>&1; then
  log "Golden instance '$GOLDEN_NAME' exists"
else
  if [ -n "${GOLDEN_TARBALL:-}" ] && [ -f "$GOLDEN_TARBALL" ]; then
    log "Importing golden from tarball $GOLDEN_TARBALL"
    run lxc import "$GOLDEN_TARBALL" "$GOLDEN_NAME" || die "golden import failed"
  else
    log "Building golden '$GOLDEN_NAME' from $BASE_IMAGE"
    run lxc launch "$BASE_IMAGE" "$GOLDEN_NAME" --profile default
    log "Waiting for agent…"
    [ "$DRY" = 1 ] || for _ in $(seq 1 60); do lxc exec "$GOLDEN_NAME" -- true 2>/dev/null && break; sleep 2; done
    log "Baking llm-station into golden"
    run "$ROOT/bin/vm-llm-station" "$GOLDEN_NAME"
  fi
  log "Stopping golden (clones use it stopped)"
  run lxc stop "$GOLDEN_NAME" 2>/dev/null || true
  run lxc snapshot "$GOLDEN_NAME" pre-clones 2>/dev/null || true
fi

# ---- 5. vm_mgr bin on PATH --------------------------------------------------
log "Linking vm_mgr/bin into /usr/local/bin"
for f in "$ROOT"/bin/vm-*; do
  [ -f "$f" ] || continue
  run ln -sf "$f" "/usr/local/bin/$(basename "$f")"
done

# ---- 6. console: venv, TLS, password, unit ----------------------------------
if [ "$SKIP_CONSOLE" = 1 ]; then
  log "Skipping console (--skip-console)"
  render_units
  log "Done."; exit 0
fi
CDIR="$ROOT/console"; VENV="$CDIR/.venv"; VPY="$VENV/bin/python3"
if [ ! -x "$VPY" ]; then
  log "Creating console venv (+ aiohttp)"
  have python3 || run apt-get install -y python3 python3-venv
  run python3 -m venv "$VENV"
  run "$VENV/bin/pip" install --quiet --upgrade pip aiohttp
fi

if [ -f "$CDIR/console.crt" ] && [ -f "$CDIR/console.key" ]; then
  log "TLS cert present"
else
  log "Generating self-signed TLS cert for $CONSOLE_DOMAIN"
  run openssl req -x509 -newkey rsa:2048 -nodes -days 3650 \
    -keyout "$CDIR/console.key" -out "$CDIR/console.crt" -subj "/CN=$CONSOLE_DOMAIN"
fi

if [ -f "$CDIR/.auth" ]; then
  log "Admin password already set (.auth present)"
elif [ "$DRY" = 1 ]; then
  echo "  [dry-run] would set admin password"
else
  PW="${ADMIN_PASSWORD:-}"
  if [ -z "$PW" ]; then
    read -r -s -p "Console admin password: " PW; echo
    [ -n "$PW" ] || die "empty password"
  fi
  log "Writing admin password hash"
  CONSOLE_PW="$PW" "$VPY" - "$CDIR/server.py" "$CDIR/.auth" <<'PY'
import importlib.util, os, sys
src, out = sys.argv[1], sys.argv[2]
spec = importlib.util.spec_from_file_location("console_server", src)
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)
open(out, "w").write(m.hash_password(os.environ["CONSOLE_PW"]))
PY
  chmod 600 "$CDIR/.auth"; chown "$CONSOLE_USER" "$CDIR/.auth"
fi

log "Writing systemd unit station-console.service"
UNIT=/etc/systemd/system/station-console.service
if [ "$DRY" = 1 ]; then
  echo "  [dry-run] would write $UNIT (User=$CONSOLE_USER ExecStart=$VPY $CDIR/server.py)"
else
  cat > "$UNIT" <<UNITEOF
[Unit]
Description=Station Console — web console for the LXD dev stations
After=network-online.target snap.lxd.daemon.service
Wants=network-online.target

[Service]
Type=simple
User=$CONSOLE_USER
Group=lxd
WorkingDirectory=$CDIR
Environment=HOST=$CONSOLE_HOST
Environment=PORT=$CONSOLE_PORT
Environment=STATION_CONSOLE_UID=$STATION_UID
Environment=STATION_CONSOLE_GID=$STATION_GID
Environment=STATION_CONSOLE_HOME=$STATION_HOME
ExecStart=$VPY $CDIR/server.py
Restart=on-failure
RestartSec=3

[Install]
WantedBy=multi-user.target
UNITEOF
fi

# ---- 7. render unit templates + enable/start the console --------------------
render_units
log "Enabling + starting station-console"
run systemctl daemon-reload
run systemctl enable --now station-console

echo
log "Done. Console: https://$CONSOLE_DOMAIN:$CONSOLE_PORT  (self-signed; admin login)"
echo "  status : systemctl status station-console"
echo "  logs   : journalctl -u station-console -f"
echo "  create : vm-new <name> [--cpu N --mem G --disk G --project P --net lan]"