#!/usr/bin/env bash
# Install / refresh the CIS490 lab-host role.
#
# Idempotent — safe to re-run after `git pull`. Does NOT enroll the
# host into WireGuard (that's wg-enroll's job, run separately and
# *first*) and does NOT mint TLS certs (that's wg-pki's job).
#
# Steps:
#   1. Verify prereqs (KVM, zstd, qemu, python3.11+, systemd).
#   2. Create the cis490 service user + /var/lib/cis490 layout.
#   3. Sync the repo into /opt/cis490 and build a uv-managed venv.
#   4. Install systemd units from etc/.
#   5. Drop /etc/cis490/lab-host.toml (only on first install).
#
# Operator finishes by:
#   - editing /etc/cis490/lab-host.toml (host_id, receiver URL, certs)
#   - placing leaf certs at /etc/cis490/certs/{lab-host.pem,key,wg-ca.pem}
#   - `systemctl enable --now cis490-shipper`

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
INSTALL_ROOT="${INSTALL_ROOT:-/opt/cis490}"
DATA_ROOT="${DATA_ROOT:-/var/lib/cis490}"
ETC_ROOT="${ETC_ROOT:-/etc/cis490}"
SERVICE_USER="${SERVICE_USER:-cis490}"

log() { printf '[install-lab-host] %s\n' "$*" >&2; }
die() { log "FATAL: $*"; exit 1; }

# --- 1. prereqs --------------------------------------------------------
log "checking prereqs"

if [[ $EUID -ne 0 ]]; then
    die "must run as root (writes to /opt, /etc, /var/lib, and systemd)"
fi
command -v systemctl >/dev/null || die "systemd not found"
command -v qemu-system-x86_64 >/dev/null || die "qemu-system-x86_64 not on PATH"
command -v zstd >/dev/null || die "zstd not on PATH (apt install zstd)"
[[ -e /dev/kvm ]] || die "/dev/kvm missing — KVM not available"

# uv is preferred (lockfile-driven). Fall back to system pip if absent.
USE_UV=0
if command -v uv >/dev/null; then USE_UV=1; fi

# --- 2. user + layout --------------------------------------------------
log "ensuring service user $SERVICE_USER"
if ! id -u "$SERVICE_USER" >/dev/null 2>&1; then
    useradd --system --no-create-home --shell /usr/sbin/nologin \
        --home-dir "$INSTALL_ROOT" "$SERVICE_USER"
fi
# kvm group lets the service spawn VMs.
if getent group kvm >/dev/null 2>&1; then
    usermod -a -G kvm "$SERVICE_USER" || true
fi

install -d -o root -g root -m 0755 "$ETC_ROOT" "$ETC_ROOT/certs"
install -d -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0755 \
    "$DATA_ROOT" "$DATA_ROOT/data" \
    "$DATA_ROOT/data/episodes" "$DATA_ROOT/data/outbox" \
    "$DATA_ROOT/data/shipped" "$DATA_ROOT/data/queue" \
    "$DATA_ROOT/samples" "$DATA_ROOT/samples/store" \
    "$DATA_ROOT/vm" "$DATA_ROOT/vm/images"

# --- 3. repo + venv ----------------------------------------------------
log "syncing repo into $INSTALL_ROOT"
install -d -o "$SERVICE_USER" -g "$SERVICE_USER" -m 0755 "$INSTALL_ROOT"
# We use a clean cp -aT rather than rsync to avoid an extra dep.
cp -aT "$REPO_ROOT" "$INSTALL_ROOT"
chown -R "$SERVICE_USER":"$SERVICE_USER" "$INSTALL_ROOT"

log "building venv"
if [[ "$USE_UV" -eq 1 ]]; then
    sudo -u "$SERVICE_USER" -- env HOME="$INSTALL_ROOT" \
        uv sync --project "$INSTALL_ROOT"
else
    sudo -u "$SERVICE_USER" -- python3 -m venv "$INSTALL_ROOT/.venv"
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/pip" install \
        --quiet --upgrade pip
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/pip" install \
        --quiet starlette 'uvicorn[standard]' httpx msgpack
fi

# --- 4. systemd --------------------------------------------------------
log "installing systemd units"
install -m 0644 "$REPO_ROOT/etc/cis490-shipper.service" \
    /etc/systemd/system/cis490-shipper.service
install -m 0644 "$REPO_ROOT/etc/cis490-orchestrator.service" \
    /etc/systemd/system/cis490-orchestrator.service
systemctl daemon-reload

# --- 5. config template (only on first install) -----------------------
if [[ ! -f "$ETC_ROOT/lab-host.toml" ]]; then
    log "writing $ETC_ROOT/lab-host.toml (template)"
    install -m 0640 -o root -g "$SERVICE_USER" \
        "$REPO_ROOT/etc/lab-host.toml.example" "$ETC_ROOT/lab-host.toml"
    NEW_INSTALL=1
else
    log "$ETC_ROOT/lab-host.toml exists; leaving in place"
    NEW_INSTALL=0
fi

# --- 6. orchestrator env file (read by cis490-orchestrator.service) ----
ENV_FILE="$ETC_ROOT/lab-host.env"
DEFAULT_HOST_ID="$(hostname -s)"
if [[ ! -f "$ENV_FILE" ]]; then
    log "writing $ENV_FILE (host_id defaults to $DEFAULT_HOST_ID — edit if you want something else)"
    install -m 0640 -o root -g "$SERVICE_USER" /dev/stdin "$ENV_FILE" <<EOF
# Read by cis490-orchestrator.service. Override per-host as needed.
FLEET_HOST_ID=$DEFAULT_HOST_ID
# BRIDGE=br-malware     # uncomment to enable source 4 pcap capture
EOF
fi

# --- 7. mTLS leaf cert (auto-fetch via bootstrap.wg) -------------------
# Pull our leaf cert from the Pi's bootstrap endpoint if it isn't
# already on disk. Trust boundary: "reached bootstrap.wg over WG"
# (iptmonads already filters non-peers from 443). Caddy's TLS cert
# is verified against the bundled etc/caddy-root.crt — no chicken-
# and-egg.
HOST_ID="$(grep -E '^host_id\s*=' "$ETC_ROOT/lab-host.toml" 2>/dev/null \
    | head -1 | sed -E 's/^host_id\s*=\s*"([^"]+)".*/\1/')"
if [[ -z "$HOST_ID" || "$HOST_ID" == "REPLACE_ME" ]]; then
    log "skipping cert auto-fetch: host_id not set in $ETC_ROOT/lab-host.toml"
elif [[ ! -f "$ETC_ROOT/certs/lab-host.pem" ]]; then
    log "fetching leaf cert from https://bootstrap.wg/v1/cert/$HOST_ID"
    install -d -m 0755 -o root -g "$SERVICE_USER" "$ETC_ROOT/certs"
    TAR="/tmp/cis490-bootstrap-$$.tar"
    if curl -fsS --cacert "$REPO_ROOT/etc/caddy-root.crt" \
            --connect-timeout 10 --max-time 60 \
            "https://bootstrap.wg/v1/cert/$HOST_ID" -o "$TAR"; then
        tar -C "$ETC_ROOT/certs" -xf "$TAR"
        mv "$ETC_ROOT/certs/ca.crt"          "$ETC_ROOT/certs/wg-ca.pem"
        mv "$ETC_ROOT/certs/$HOST_ID.pem"    "$ETC_ROOT/certs/lab-host.pem"
        mv "$ETC_ROOT/certs/$HOST_ID.key"    "$ETC_ROOT/certs/lab-host.key"
        chown root:"$SERVICE_USER" "$ETC_ROOT/certs/"*.pem \
            "$ETC_ROOT/certs/lab-host.key"
        chmod 0644 "$ETC_ROOT/certs/"*.pem
        chmod 0640 "$ETC_ROOT/certs/lab-host.key"
        rm -f "$TAR"
        log "leaf cert installed for host_id=$HOST_ID"
    else
        rm -f "$TAR"
        log "WARN: bootstrap.wg fetch failed — make sure /etc/hosts maps it"
        log "      to 10.100.0.1 and that wg0 is up. cert delivery skipped."
    fi
else
    log "$ETC_ROOT/certs/lab-host.pem present; skipping auto-fetch"
fi

# --- 8. baseline VM image + cidata (best-effort) -----------------------
ALPINE_IMG="$DATA_ROOT/vm/images/alpine-baseline.qcow2"
CIDATA_ISO="$DATA_ROOT/vm/images/cidata.iso"
if [[ ! -f "$ALPINE_IMG" ]]; then
    if "$REPO_ROOT/scripts/fetch-alpine-baseline.sh" "$ALPINE_IMG"; then
        log "fetched Alpine baseline -> $ALPINE_IMG"
    else
        log "WARN: Alpine baseline fetch failed; drop a qcow2 at $ALPINE_IMG manually"
    fi
fi
if [[ -f "$ALPINE_IMG" && ! -f "$CIDATA_ISO" ]]; then
    log "building cidata.iso (in-guest agent embedded)"
    sudo -u "$SERVICE_USER" -- "$INSTALL_ROOT/.venv/bin/python" \
        "$INSTALL_ROOT/tools/build_cidata.py" "$CIDATA_ISO" || \
        log "WARN: cidata build failed; run tools/build_cidata.py manually"
fi
# Symlink the canonical paths the launchers look at, when missing.
ln -sf "$ALPINE_IMG" "$INSTALL_ROOT/vm/images/alpine-baseline.qcow2" 2>/dev/null || true
ln -sf "$CIDATA_ISO" "$INSTALL_ROOT/vm/images/cidata.iso" 2>/dev/null || true

if [[ "$NEW_INSTALL" == "1" ]]; then
    log ""
    log "================================================================="
    log "  FIRST-INSTALL NEXT STEPS                                        "
    log "================================================================="
    log "  1. Edit $ETC_ROOT/lab-host.toml — set host_id and receiver URL."
    log ""
    log "  2. (On the Pi.) Mint + ship a leaf cert for this host:"
    log "       sudo wg-pki/scripts/deploy-cis490-cert.sh <host_id> <wg_ip>"
    log ""
    log "  3. Run the diagnostic — every red row prints the exact fix:"
    log "       $INSTALL_ROOT/.venv/bin/python \\"
    log "           $INSTALL_ROOT/tools/cis490_doctor.py --role lab-host"
    log ""
    log "  4. Smoke-test the pipe (returns ok=true on success):"
    log "       sudo -u $SERVICE_USER $INSTALL_ROOT/.venv/bin/python -m shipper \\"
    log "            --config $ETC_ROOT/lab-host.toml --ping"
    log ""
    log "  5. Turn on the services — episodes start flowing immediately:"
    log "       sudo systemctl enable --now cis490-shipper cis490-orchestrator"
    log "================================================================="
fi

log "lab-host install complete."
log ""
log "Cloning this repo and running the launchers manually is NOT enough."
log "The lab-host role's data flow lives in the systemd services this"
log "script just installed. If $INSTALL_ROOT/index.jsonl on the Pi stays"
log "empty after step 5, run:"
log "   $INSTALL_ROOT/.venv/bin/python $INSTALL_ROOT/tools/cis490_doctor.py"