#!/usr/bin/env bash
# WAFio All-in-One Installer — Free Plan
# https://wafio.cloud/docs#installation
#
# Usage:
#   curl -fsSL https://releases.wafio.cloud/install.sh | sudo bash
#   sudo bash install.sh
#   sudo bash install.sh --component controlplane
#
# This script handles everything end-to-end:
#   1. License key input + hardware binding (wafio.cloud portal)
#   2. Component selection (controlplane / waf-agent / host-agent / all)
#   3. OS detection (Debian/Ubuntu or RHEL/CentOS/Rocky)
#   4. PostgreSQL check / auto-install
#   5. Binary download from releases.wafio.cloud
#   6. Security data download (GeoIP, CRS, FireHOL, JA3, Bots)
#   7. Config file generation (/etc/wafio/config.yaml)
#   8. Database migrations
#   9. Admin account creation (no browser wizard needed)
#  11. Project + agent credential provisioning
#  12. Systemd service installation + start
#  13. Health verification
#
# Usage:
#   curl -fsSL https://releases.wafio.cloud/install.sh | sudo bash
#   curl -fsSL https://releases.wafio.cloud/install.sh | sudo bash -s -- controlplane
#   curl -fsSL https://releases.wafio.cloud/install.sh | sudo bash -s -- waf-agent
#   curl -fsSL https://releases.wafio.cloud/install.sh | sudo bash -s -- host-agent
#
# Or download first, then run (recommended for production):
#   curl -fsSL https://releases.wafio.cloud/install.sh -o wafio-install.sh
#   sudo bash wafio-install.sh [controlplane|waf-agent|host-agent]

set -euo pipefail

# ─── Color helpers ───────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; BOLD='\033[1m'; GRAY='\033[0;90m'; NC='\033[0m'
info()   { echo -e "${BLUE}[INFO]${NC}   $*"; }
ok()     { echo -e "${GREEN}[ ✓ ]${NC}   $*"; }
warn()   { echo -e "${YELLOW}[WARN]${NC}  $*"; }
die()    { echo -e "\n${RED}[ERR]${NC}   $*" >&2; exit 1; }
step()   { echo -e "\n${BOLD}${CYAN}━━ $* ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"; }
prompt() { echo -en "  ${BOLD}→ $*${NC}"; }
hr()     { echo -e "${GRAY}──────────────────────────────────────────────────────${NC}"; }

# ─── Banner ──────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}${CYAN}"
cat <<'BANNER'
  ██╗    ██╗ █████╗ ███████╗██╗ ██████╗
  ██║    ██║██╔══██╗██╔════╝██║██╔═══██╗
  ██║ █╗ ██║███████║█████╗  ██║██║   ██║
  ██║███╗██║██╔══██║██╔══╝  ██║██║   ██║
  ╚███╔███╔╝██║  ██║██║     ██║╚██████╔╝
   ╚══╝╚══╝ ╚═╝  ╚═╝╚═╝     ╚═╝ ╚═════╝
BANNER
echo -e "${NC}"
echo -e "  ${BOLD}WAFio — Web Application Firewall Platform${NC}"
echo -e "  ${GRAY}Free Plan Installer v2.0 — All-in-One Setup${NC}"
echo ""
hr

# ─── Prerequisites ───────────────────────────────────────────────────────────
[[ $EUID -eq 0 ]] || die "Run as root: sudo bash $0"
command -v curl    &>/dev/null || die "curl required: apt-get install curl"
command -v tar     &>/dev/null || die "tar required"
command -v openssl &>/dev/null || die "openssl required"
command -v python3 &>/dev/null || die "python3 required"

# ─── Defaults ────────────────────────────────────────────────────────────────
PORTAL_API_URL="${PORTAL_API_URL:-https://wafio.cloud}"
RELEASES_URL="${RELEASES_URL:-https://releases.wafio.cloud}"
INSTALL_DIR="/opt/wafio"
CONFIG_DIR="/etc/wafio"
LOG_DIR="/var/log/wafio"
DATA_DIR="/var/lib/wafio"
SECURITY_DATA_DIR="${INSTALL_DIR}/security-data"
CONFIG_FILE="${CONFIG_DIR}/config.yaml"
CERTS_DIR="${CONFIG_DIR}/certs"
CP_HTTP_PORT=9087
CP_GRPC_PORT=9090
COMPONENT=""
ADMIN_EMAIL=""; ADMIN_PASS=""; ADMIN_NAME=""; PROJECT_ID=""; PROJECT_NAME=""
GRPC_DOMAIN=""; JWT=""
PG_HOST="localhost"; PG_PORT="5432"; PG_USER="postgres"; PG_PASS=""; PG_DB="wafio"

# ─── Parse flags ─────────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
  case "$1" in
    --component) COMPONENT="$2"; shift 2 ;;
    controlplane|waf-agent|host-agent|all) COMPONENT="$1"; shift ;;
    *) shift ;;
  esac
done

# ─── 1. Component selection ──────────────────────────────────────────────────
step "Component Selection"
if [[ -z "$COMPONENT" ]]; then
  echo ""
  echo "  What would you like to install?"
  echo ""
  echo "    1) All Components (recommended)  — Control Plane + WAF Agent + Host Agent"
  echo "    2) Control Plane only            — WAFio server + web dashboard"
  echo "    3) WAF Agent only                — HTTP L7 inspection sidecar"
  echo "    4) Host Agent only               — eBPF kernel firewall (Linux 5.8+)"
  echo ""
  prompt "Your choice [1-4]: "; read -r CHOICE
  case "$CHOICE" in
    1) COMPONENT="all" ;;
    2) COMPONENT="controlplane" ;;
    3) COMPONENT="waf-agent" ;;
    4) COMPONENT="host-agent" ;;
    *) die "Invalid choice: $CHOICE" ;;
  esac
fi

INSTALL_CP=false; INSTALL_WAF=false; INSTALL_HOST=false
case "$COMPONENT" in
  all)          INSTALL_CP=true;  INSTALL_WAF=true;  INSTALL_HOST=true ;;
  controlplane) INSTALL_CP=true ;;
  waf-agent)    INSTALL_WAF=true ;;
  host-agent)   INSTALL_HOST=true ;;
  *) die "Unknown component: $COMPONENT" ;;
esac

echo ""
echo -e "  Will install: ${BOLD}"
$INSTALL_CP   && echo "    • Control Plane"
$INSTALL_WAF  && echo "    • WAF Agent"
$INSTALL_HOST && echo "    • Host Agent"
echo -e "${NC}"

# ─── 2. License key input ────────────────────────────────────────────────────
step "License Verification"
echo ""
echo "  A license key is required from the WAFio portal."
echo "  → Register free at: ${PORTAL_API_URL}/register"
echo "  → After login: Dashboard → Generate License → copy the token"
echo ""

declare -A LICENSE_TOKENS=()
for COMP in controlplane waf-agent host-agent; do
  SHOULD_INSTALL=false
  case "$COMP" in
    controlplane) $INSTALL_CP   && SHOULD_INSTALL=true ;;
    waf-agent)    $INSTALL_WAF  && SHOULD_INSTALL=true ;;
    host-agent)   $INSTALL_HOST && SHOULD_INSTALL=true ;;
  esac
  if $SHOULD_INSTALL; then
    while true; do
      prompt "License token for ${COMP}: "; read -r TK; TK="${TK// /}"
      [[ -n "$TK" ]] && { LICENSE_TOKENS["$COMP"]="$TK"; break; }
      echo "  ✗ Token cannot be empty."
    done
  fi
done

# ─── 3. Hardware ID (matches internal/license/hardware.go) ───────────────────
step "Hardware Binding"
if [[ -n "${WAFIO_MACHINE_ID:-}" ]]; then
  HARDWARE_ID=$(printf "env:%s" "$WAFIO_MACHINE_ID" | sha256sum | awk '{print $1}')
elif [[ -f /etc/machine-id ]]; then
  _mid=$(tr -d '[:space:]' < /etc/machine-id)
  HARDWARE_ID=$(printf "machine-id:%s" "$_mid" | sha256sum | awk '{print $1}')
else
  _macs=$(ip link show 2>/dev/null | awk '/link\/ether/{print $2}' | sort | tr '\n' '|' | sed 's/|$//')
  HARDWARE_ID=$(printf "mac:%s|[%s]" "$(hostname)" "$_macs" | sha256sum | awk '{print $1}')
fi
HOSTNAME_VAL=$(hostname -f 2>/dev/null || hostname)
ok "Hardware fingerprint bound to: ${HOSTNAME_VAL}"

# ─── 4. Activate all licenses via portal API ─────────────────────────────────
step "Activating Licenses"
declare -A LICENSE_JSONS=()

for COMP in "${!LICENSE_TOKENS[@]}"; do
  info "Activating license for ${COMP}..."
  TK="${LICENSE_TOKENS[$COMP]}"
  _RESP_FILE=$(mktemp)
  HTTP_CODE=$(curl -s --max-time 20 -o "$_RESP_FILE" -w "%{http_code}" \
    -X POST "${PORTAL_API_URL}/api/license/activate" \
    -H "Content-Type: application/json" \
    -d "{\"token\":\"${TK}\",\"hardware_id\":\"${HARDWARE_ID}\",\"hostname\":\"${HOSTNAME_VAL}\"}" 2>/dev/null) \
    || die "Cannot reach ${PORTAL_API_URL}. Check internet connectivity."
  RESP=$(cat "$_RESP_FILE"); rm -f "$_RESP_FILE"

  if [[ "$HTTP_CODE" != "200" ]]; then
    ERR=$(echo "$RESP" | python3 -c "import sys,json;d=json.load(sys.stdin);print(d.get('error','unknown'))" 2>/dev/null || echo "$RESP")
    die "License activation failed for ${COMP} (HTTP ${HTTP_CODE}): ${ERR}"
  fi

  LIC_JSON=$(echo "$RESP" | python3 -c "
import sys, json
try:
    resp = json.load(sys.stdin)
    val = resp.get('license', '')
    if isinstance(val, str):
        print(json.dumps(json.loads(val)))
    elif isinstance(val, dict):
        print(json.dumps(val))
    else:
        sys.exit(1)
except Exception:
    sys.exit(1)
" 2>/dev/null || true)

  if [[ -z "$LIC_JSON" ]]; then
    die "License activation failed for ${COMP}: empty license in response"
  fi
  LICENSE_JSONS["$COMP"]="$LIC_JSON"
  ok "${COMP} license activated and bound to this hardware"
done

# ─── 5. OS detection ─────────────────────────────────────────────────────────
step "Operating System Detection"
OS_FAMILY="unknown"; OS_MAJOR=8; PRETTY_NAME="unknown"
if [[ -f /etc/os-release ]]; then
  . /etc/os-release
  case "${ID:-}" in
    ubuntu|debian|raspbian|linuxmint) OS_FAMILY="debian" ;;
    centos|rhel|fedora|rocky|almalinux|ol)
      OS_FAMILY="rhel"
      OS_MAJOR=$(echo "${VERSION_ID:-8}" | cut -d. -f1) ;;
    *) warn "Unrecognized OS: ${ID:-unknown}. Dependency auto-install disabled." ;;
  esac
fi
ok "OS: ${PRETTY_NAME:-unknown} (family: ${OS_FAMILY})"

pkg_install() {
  case "$OS_FAMILY" in
    debian) DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends "$@" ;;
    rhel)   yum install -y "$@" 2>/dev/null || dnf install -y "$@" ;;
    *) warn "Cannot auto-install: $*. Install manually." ;;
  esac
}
pkg_update() {
  case "$OS_FAMILY" in
    debian) DEBIAN_FRONTEND=noninteractive apt-get update -qq ;;
    rhel)   yum makecache -q 2>/dev/null || dnf makecache -q ;;
    *) : ;;
  esac
}

# ─── 6. PostgreSQL (control plane only) ──────────────────────────────────────
if $INSTALL_CP; then
  step "Database Setup (PostgreSQL)"
  echo ""
  echo "  WAFio Control Plane requires PostgreSQL 17+."
  echo ""
  prompt "Is PostgreSQL already installed and running? [y/N]: "; read -r PG_EXISTS
  PG_EXISTS="${PG_EXISTS:-n}"

  if [[ "$PG_EXISTS" =~ ^[Yy] ]]; then
    prompt "PostgreSQL host [localhost]: ";  read -r PG_HOST; PG_HOST="${PG_HOST:-localhost}"
    prompt "PostgreSQL port [5432]: ";       read -r PG_PORT; PG_PORT="${PG_PORT:-5432}"
    prompt "PostgreSQL admin user [postgres]: "; read -r PG_ADMIN; PG_ADMIN="${PG_ADMIN:-postgres}"
    prompt "PostgreSQL admin password (leave empty for peer auth): "; read -rs PG_ADMIN_PASS; echo ""
    prompt "WAFio database name [wafio]: ";  read -r PG_DB;   PG_DB="${PG_DB:-wafio}"

    info "Testing PostgreSQL connection..."
    # Test connectivity using peer auth or password
    if [[ -z "$PG_ADMIN_PASS" && "$PG_HOST" == "localhost" ]]; then
      sudo -u postgres psql -p "$PG_PORT" -c '\q' &>/dev/null \
        || die "Cannot connect to PostgreSQL via peer auth. Try setting an admin password."
    else
      PGPASSWORD="$PG_ADMIN_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_ADMIN" -c '\q' &>/dev/null \
        || die "Cannot connect to PostgreSQL. Verify credentials and retry."
    fi

    # Create dedicated wafio user + db
    PG_USER="wafio"; PG_PASS=$(openssl rand -hex 16)
    info "Creating database user 'wafio' and database '${PG_DB}'..."
    if [[ -z "$PG_ADMIN_PASS" && "$PG_HOST" == "localhost" ]]; then
      # Always create (ignore error if exists), then ALWAYS force-set password
      sudo -u postgres psql -p "$PG_PORT" -c "CREATE USER wafio;" 2>/dev/null || true
      sudo -u postgres psql -p "$PG_PORT" -c "ALTER USER wafio WITH PASSWORD '${PG_PASS}';" \
        || die "Failed to set password for PostgreSQL user 'wafio'"
      sudo -u postgres psql -p "$PG_PORT" -c "CREATE DATABASE ${PG_DB} OWNER wafio;" 2>/dev/null || true
      sudo -u postgres psql -p "$PG_PORT" -c "GRANT ALL ON DATABASE ${PG_DB} TO wafio;" 2>/dev/null || true
    else
      PGPASSWORD="$PG_ADMIN_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_ADMIN" \
        -c "CREATE USER wafio;" 2>/dev/null || true
      PGPASSWORD="$PG_ADMIN_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_ADMIN" \
        -c "ALTER USER wafio WITH PASSWORD '${PG_PASS}';" \
        || die "Failed to set password for PostgreSQL user 'wafio'"
      PGPASSWORD="$PG_ADMIN_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_ADMIN" \
        -c "CREATE DATABASE ${PG_DB} OWNER wafio;" 2>/dev/null || true
      PGPASSWORD="$PG_ADMIN_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_ADMIN" \
        -c "GRANT ALL ON DATABASE ${PG_DB} TO wafio;" 2>/dev/null || true
    fi
    ok "Database '${PG_DB}' ready"
  else
    info "Installing PostgreSQL 17..."
    case "$OS_FAMILY" in
      debian)
        pkg_update && pkg_install gnupg2 lsb-release ca-certificates
        curl -fsSL https://www.postgresql.org/media/keys/ACCC4CF8.asc \
          | gpg --dearmor -o /usr/share/keyrings/postgresql.gpg
        echo "deb [signed-by=/usr/share/keyrings/postgresql.gpg] https://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" \
          > /etc/apt/sources.list.d/pgdg.list
        pkg_update && pkg_install postgresql-17
        systemctl enable --now postgresql
        ;;
      rhel)
        PGDG="https://download.postgresql.org/pub/repos/yum/reporpms/EL-${OS_MAJOR}-x86_64/pgdg-redhat-repo-latest.noarch.rpm"
        yum install -y "$PGDG" 2>/dev/null || dnf install -y "$PGDG"
        yum install -y postgresql17-server 2>/dev/null || dnf install -y postgresql17-server
        /usr/pgsql-17/bin/postgresql-17-setup initdb 2>/dev/null || true
        systemctl enable --now postgresql-17
        ;;
      *) die "Please install PostgreSQL 17+ manually, then re-run." ;;
    esac
    sleep 2
    PG_HOST="localhost"; PG_PORT="5432"; PG_USER="wafio"
    PG_PASS=$(openssl rand -hex 16); PG_DB="wafio"
    info "Creating database user '${PG_USER}' and database '${PG_DB}'..."
    sudo -u postgres psql -c "CREATE USER ${PG_USER} WITH PASSWORD '${PG_PASS}';" 2>/dev/null || true
    sudo -u postgres psql -c "CREATE DATABASE ${PG_DB} OWNER ${PG_USER};" 2>/dev/null || true
    sudo -u postgres psql -c "GRANT ALL ON DATABASE ${PG_DB} TO ${PG_USER};" 2>/dev/null || true
    ok "Database '${PG_DB}' created"
  fi

  info "Testing database connection..."
  PGPASSWORD="$PG_PASS" psql -h "$PG_HOST" -p "$PG_PORT" -U "$PG_USER" -d "$PG_DB" -c '\q' &>/dev/null \
    || die "Cannot connect to PostgreSQL. Verify credentials and retry."
  ok "PostgreSQL connection verified"

  # ─── TimescaleDB ─────────────────────────────────────────────────────────────
  step "TimescaleDB Setup"
  _tsdb_ok=false
  if sudo -u postgres psql -p "$PG_PORT" -d "$PG_DB" -c '\dx timescaledb' 2>/dev/null | grep -q timescaledb; then
    ok "TimescaleDB extension already enabled"
    _tsdb_ok=true
  fi
  if ! $_tsdb_ok; then
    info "Installing TimescaleDB for PostgreSQL 17..."
    case "$OS_FAMILY" in
      debian)
        _CODENAME=$(lsb_release -cs 2>/dev/null || echo "jammy")
        curl -fsSL https://packagecloud.io/timescale/timescaledb/gpgkey \
          | gpg --dearmor -o /usr/share/keyrings/timescaledb.gpg 2>/dev/null
        echo "deb [signed-by=/usr/share/keyrings/timescaledb.gpg] https://packagecloud.io/timescale/timescaledb/ubuntu/ ${_CODENAME} main" \
          > /etc/apt/sources.list.d/timescaledb.list
        pkg_update
        pkg_install timescaledb-2-postgresql-17 \
          || die "TimescaleDB install failed. See: https://docs.timescale.com/self-hosted/latest/install/"
        # Preload timescaledb — required before CREATE EXTENSION
        _PG_CONF="/etc/postgresql/17/main/postgresql.conf"
        if [[ -f "$_PG_CONF" ]]; then
          if ! grep -q "timescaledb" "$_PG_CONF"; then
            echo "shared_preload_libraries = 'timescaledb'" >> "$_PG_CONF"
          fi
        fi
        timescaledb-tune --quiet --yes 2>/dev/null || true
        systemctl restart postgresql
        sleep 3
        ;;
      rhel)
        yum install -y timescaledb-2-postgresql-17 2>/dev/null \
          || dnf install -y timescaledb-2-postgresql-17 2>/dev/null \
          || die "TimescaleDB install failed. See: https://docs.timescale.com/self-hosted/latest/install/rhel/"
        _PG_CONF="/var/lib/pgsql/17/data/postgresql.conf"
        if [[ -f "$_PG_CONF" ]]; then
          if ! grep -q "timescaledb" "$_PG_CONF"; then
            echo "shared_preload_libraries = 'timescaledb'" >> "$_PG_CONF"
          fi
        fi
        timescaledb-tune --quiet --yes 2>/dev/null || true
        systemctl restart postgresql-17 2>/dev/null || systemctl restart postgresql
        sleep 3
        ;;
      *) die "Cannot auto-install TimescaleDB on this OS. See: https://docs.timescale.com/self-hosted/latest/install/" ;;
    esac
    sudo -u postgres psql -p "$PG_PORT" -d "$PG_DB" \
      -c "CREATE EXTENSION IF NOT EXISTS timescaledb;" \
      && ok "TimescaleDB installed and enabled" \
      || die "Could not enable TimescaleDB extension. Run manually: sudo -u postgres psql ${PG_DB} -c 'CREATE EXTENSION timescaledb;'"
  fi

fi

# ─── 7. Download binaries ─────────────────────────────────────────────────────
step "Downloading WAFio Binaries"
LATEST_JSON=$(curl -fsSL --max-time 15 "${RELEASES_URL}/latest.json" 2>/dev/null || echo '{"version":"dev"}')
VERSION=$(echo "$LATEST_JSON" | python3 -c "import sys,json;print(json.load(sys.stdin).get('version','dev'))" 2>/dev/null || echo "dev")
info "WAFio version: ${VERSION}"

mkdir -p "$INSTALL_DIR" "$CONFIG_DIR" "${CERTS_DIR}/waf" "${CERTS_DIR}/host" \
  "$LOG_DIR" "$DATA_DIR" "${INSTALL_DIR}/migrations"

download_component() {
  local COMP="$1" BINARY_NAME="$2" DEST="$3"
  local BUNDLE="wafio-${COMP}-${VERSION}-linux-amd64.tar.gz"
  local URL="${RELEASES_URL}/bundles/${BUNDLE}"
  local TMP; TMP=$(mktemp -d)

  # Check for pre-staged bundle in /tmp (bypasses download for offline/slow installs)
  local STAGED="/tmp/${BUNDLE}"
  if [[ -f "$STAGED" ]]; then
    info "Using pre-staged bundle: ${STAGED}"
    cp "$STAGED" "${TMP}/${BUNDLE}"
  else
    info "Downloading ${BUNDLE}..."
    local _tries=0 _ok=false
    while [[ $_tries -lt 3 ]]; do
      _tries=$((_tries+1))
      curl -fL --retry 3 --retry-delay 5 --retry-max-time 120 \
        -C - --progress-bar "${URL}" -o "${TMP}/${BUNDLE}" && { _ok=true; break; }
      warn "Download attempt ${_tries}/3 failed, retrying..."
      sleep 5
    done
    if ! $_ok; then rm -rf "$TMP"; die "Failed to download ${URL} after 3 attempts"; fi
  fi

  tar xzf "${TMP}/${BUNDLE}" -C "${TMP}" 2>/dev/null || true
  # Filter macOS metadata entries (._*) to find the real top-level directory
  local EXDIR="${TMP}/$(tar tzf "${TMP}/${BUNDLE}" 2>/dev/null | grep -v '^\._\|/\._' | head -1 | cut -d/ -f1)"
  install -m 0755 "${EXDIR}/${BINARY_NAME}" "${DEST}"
  ok "Installed ${DEST}"
  if [[ "$COMP" == "controlplane" ]]; then
    [[ -d "${EXDIR}/frontend" ]]   && { rm -rf "${INSTALL_DIR}/frontend";   cp -r "${EXDIR}/frontend"   "${INSTALL_DIR}/frontend";   ok "Frontend assets installed"; }
    [[ -d "${EXDIR}/migrations" ]] && { rm -rf "${INSTALL_DIR}/migrations"; cp -r "${EXDIR}/migrations" "${INSTALL_DIR}/migrations"; ok "Migrations directory installed"; }
  fi
  if [[ "$COMP" == "host-agent" ]]; then
    # Extract eBPF objects (CO-RE compiled, pre-built) from bundle
    if [[ -d "${EXDIR}/ebpf" ]]; then
      mkdir -p "${INSTALL_DIR}/ebpf"
      cp "${EXDIR}/ebpf/"*.o "${INSTALL_DIR}/ebpf/" 2>/dev/null || true
      ok "eBPF objects installed to ${INSTALL_DIR}/ebpf/"
    fi
  fi
  rm -rf "$TMP"
}

$INSTALL_CP   && download_component "controlplane" "wafio"            "/usr/local/bin/wafio"
$INSTALL_WAF  && download_component "waf-agent"    "wafio-waf-agent"  "/usr/local/bin/wafio-waf-agent"
$INSTALL_HOST && download_component "host-agent"   "wafio-host-agent" "/usr/local/bin/wafio-host-agent"

# ─── Save license files ───────────────────────────────────────────────────────
step "Installing License Files"
mkdir -p "${CONFIG_DIR}/waf-agent" "${CONFIG_DIR}/host-agent"
for COMP in "${!LICENSE_JSONS[@]}"; do
  case "$COMP" in
    controlplane)
      echo "${LICENSE_JSONS[$COMP]}" > "${CONFIG_DIR}/license.wafio"
      chmod 600 "${CONFIG_DIR}/license.wafio"
      ok "Control plane license → ${CONFIG_DIR}/license.wafio"
      ;;
    waf-agent)
      echo "${LICENSE_JSONS[$COMP]}" > "${CONFIG_DIR}/waf-agent/license.json"
      chmod 600 "${CONFIG_DIR}/waf-agent/license.json"
      ok "WAF agent license → ${CONFIG_DIR}/waf-agent/license.json"
      ;;
    host-agent)
      echo "${LICENSE_JSONS[$COMP]}" > "${CONFIG_DIR}/host-agent/license.json"
      chmod 600 "${CONFIG_DIR}/host-agent/license.json"
      ok "Host agent license → ${CONFIG_DIR}/host-agent/license.json"
      ;;
  esac
done

# ─── 8. Security data ─────────────────────────────────────────────────────────
step "Downloading Security Data"
info "Fetching: GeoIP, OWASP CRS, FireHOL, JA3 fingerprints, Bot signatures..."
info "This may take 2-5 minutes depending on connection speed."
echo ""
FETCH_PROFILE="all"
$INSTALL_CP && ! $INSTALL_WAF && ! $INSTALL_HOST && FETCH_PROFILE="control-plane"
! $INSTALL_CP && { $INSTALL_WAF || $INSTALL_HOST; } && FETCH_PROFILE="agent"

if command -v /usr/local/bin/wafio &>/dev/null; then
  (cd "${CONFIG_DIR}" && /usr/local/bin/wafio fetch-data --dir "${SECURITY_DATA_DIR}" --profile "${FETCH_PROFILE}") 2>/dev/null \
    || warn "Security data download incomplete. Re-run: wafio fetch-data --profile ${FETCH_PROFILE}"
  ok "Security data stored at ${SECURITY_DATA_DIR}"
else
  warn "Binary not yet available — run: wafio fetch-data --profile ${FETCH_PROFILE}"
fi

# ─── 9. Config + migrations + admin setup (control plane) ────────────────────
if $INSTALL_CP; then
  step "Configuring Control Plane"
  JWT_SECRET=$(openssl rand -hex 32)
  GRPC_DOMAIN=$(hostname -f 2>/dev/null || hostname)

  cat > "${CONFIG_FILE}" <<YAML_EOF
server:
  env: production
  http_addr: ":${CP_HTTP_PORT}"
  agent_grpc_addr: ":${CP_GRPC_PORT}"
  grpc_domain: "${GRPC_DOMAIN}"
  frontend_path: "${INSTALL_DIR}/frontend"

database:
  host: "${PG_HOST}"
  port: ${PG_PORT}
  user: "${PG_USER}"
  password: "${PG_PASS}"
  name: "${PG_DB}"
  sslmode: disable

jwt:
  secret: "${JWT_SECRET}"
  expire_hours: 720

license:
  key_file: "${CONFIG_DIR}/license.wafio"

geo:
  country_db_path: "${SECURITY_DATA_DIR}/geoip/country.mmdb"
  city_db_path: "${SECURITY_DATA_DIR}/geoip/city.mmdb"
  asn_db_path: "${SECURITY_DATA_DIR}/geoip/asn.mmdb"

bot:
  ua_bad_bot_file: "${SECURITY_DATA_DIR}/user-agents/bot.txt"
  ja3_blacklist_file: "${SECURITY_DATA_DIR}/ja3/ja3.txt"

crs:
  rules_dir: "${SECURITY_DATA_DIR}/crs/coreruleset"

mtls:
  certs_dir: "${CERTS_DIR}/agent"
  cert_validity_days: 365
YAML_EOF
  chmod 600 "${CONFIG_FILE}"
  # Create symlink so wafio can find config.yaml when run from INSTALL_DIR (where migrations/ lives)
  ln -sf "${CONFIG_FILE}" "${INSTALL_DIR}/config.yaml"
  ok "Config written to ${CONFIG_FILE}"

  step "Running Database Migrations"
  MIGRATIONS_PATH="${INSTALL_DIR}/migrations"
  [[ -d "$MIGRATIONS_PATH" ]] || die "Migrations not found at ${MIGRATIONS_PATH}. Check bundle integrity."
  info "Applying migrations to ${PG_DB}..."
  # Run from INSTALL_DIR: viper finds config.yaml (symlink) + migrations/ in same dir
  (cd "${INSTALL_DIR}" && /usr/local/bin/wafio migrate up) \
    || die "Migration failed. Check DB credentials in ${CONFIG_FILE}."
  ok "All database migrations applied"

  step "Creating Admin Account"
  echo ""
  info "Starting WAFio temporarily for initial setup..."
  (cd "${INSTALL_DIR}" && /usr/local/bin/wafio serve > /tmp/wafio-setup.log 2>&1) &
  SETUP_PID=$!

  # Wait for ready
  READY=false
  for i in $(seq 1 40); do
    if curl -sf "http://127.0.0.1:${CP_HTTP_PORT}/health" &>/dev/null 2>&1 \
       || curl -sf "http://127.0.0.1:${CP_HTTP_PORT}/api/setup/status" &>/dev/null 2>&1; then
      READY=true; break
    fi
    sleep 1
  done
  if ! $READY; then
    echo "--- Last 20 lines of /tmp/wafio-setup.log ---" >&2
    tail -20 /tmp/wafio-setup.log >&2 || true
    # Terminate using python to avoid bash kill variable restriction
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Control plane did not start within 40s. See /tmp/wafio-setup.log"
  fi
  ok "Control plane is ready for setup"

  echo ""
  echo "  Create your WAFio administrator account:"
  echo ""
  prompt "Full name: ";                          read -r ADMIN_NAME
  prompt "Email address: ";                      read -r ADMIN_EMAIL
  while true; do
    prompt "Password (min 8 chars): ";           read -rs ADMIN_PASS;  echo ""
    prompt "Confirm password: ";                 read -rs ADMIN_PASS2; echo ""
    if [[ "$ADMIN_PASS" == "$ADMIN_PASS2" ]] && [[ ${#ADMIN_PASS} -ge 8 ]]; then
      break
    fi
    echo -e "  ${RED}✗ Passwords do not match or too short. Try again.${NC}"
  done

  SETUP_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/setup/complete" \
    -H "Content-Type: application/json" \
    -d "{\"email\":\"${ADMIN_EMAIL}\",\"name\":\"${ADMIN_NAME}\",\"password\":\"${ADMIN_PASS}\"}" 2>&1) || {
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Setup complete failed: ${SETUP_RESP}\nSee /tmp/wafio-setup.log"
  }
  SETUP_OK=$(echo "$SETUP_RESP" | python3 -c \
    "import sys,json;d=json.load(sys.stdin);print('yes' if (d.get('ok') or d.get('user_id')) else 'no')" 2>/dev/null || echo "no")
  if [[ "$SETUP_OK" != "yes" ]]; then
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Admin creation failed: ${SETUP_RESP}"
  fi
  ok "Admin account created: ${ADMIN_EMAIL}"

  info "Logging in to provision project and agent credentials..."
  LOGIN_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/login" \
    -H "Content-Type: application/json" \
    -d "{\"email\":\"${ADMIN_EMAIL}\",\"password\":\"${ADMIN_PASS}\"}" 2>&1) || {
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Login failed: ${LOGIN_RESP}"
  }
  JWT=$(echo "$LOGIN_RESP" | python3 -c "import sys,json;print(json.load(sys.stdin)['token'])" 2>/dev/null) || {
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Could not extract JWT from login response"
  }

  echo ""
  prompt "Project name [My Infrastructure]: "; read -r PROJECT_NAME
  PROJECT_NAME="${PROJECT_NAME:-My Infrastructure}"

  PROJECT_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/projects" \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer ${JWT}" \
    -d "{\"name\":\"${PROJECT_NAME}\"}" 2>&1) || {
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Create project failed: ${PROJECT_RESP}"
  }
  PROJECT_ID=$(echo "$PROJECT_RESP" | python3 -c "import sys,json;print(json.load(sys.stdin)['id'])" 2>/dev/null) || {
    python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
    die "Could not extract project ID"
  }
  ok "Project '${PROJECT_NAME}' created (ID: ${PROJECT_ID})"

  if $INSTALL_WAF; then
    info "Provisioning WAF agent gRPC credentials..."
    CRED_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/projects/${PROJECT_ID}/agent-grpc-credentials" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer ${JWT}" \
      -d '{"name":"waf-agent-1"}' 2>&1) || { warn "WAF credential failed — provision via dashboard later"; CRED_RESP=""; }
    if [[ -n "$CRED_RESP" ]]; then
      WAF_CA=$(echo "$CRED_RESP"   | python3 -c "import sys,json;print(json.load(sys.stdin).get('ca_cert',''))"     2>/dev/null || true)
      WAF_CERT=$(echo "$CRED_RESP" | python3 -c "import sys,json;print(json.load(sys.stdin).get('client_cert',''))" 2>/dev/null || true)
      WAF_KEY=$(echo "$CRED_RESP"  | python3 -c "import sys,json;print(json.load(sys.stdin).get('client_key',''))"  2>/dev/null || true)
      if [[ -n "$WAF_CA" ]] && [[ -n "$WAF_KEY" ]]; then
        echo "$WAF_CA"   > "${CERTS_DIR}/waf/ca.crt"
        echo "$WAF_CERT" > "${CERTS_DIR}/waf/client.crt"
        echo "$WAF_KEY"  > "${CERTS_DIR}/waf/client.key"
        chmod 600 "${CERTS_DIR}/waf/client.key"
        ok "WAF agent certificates saved to ${CERTS_DIR}/waf/"
      fi
    fi
  fi

  if $INSTALL_HOST; then
    info "Registering host and provisioning host agent credentials..."
    HOST_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/projects/${PROJECT_ID}/hosts" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer ${JWT}" \
      -d "{\"name\":\"$(hostname)\",\"description\":\"Provisioned by install.sh\"}" 2>&1) || HOST_RESP=""
    H_KEY_VAL=$(echo "$HOST_RESP" | python3 -c "import sys,json;print(json.load(sys.stdin).get('host_key',''))" 2>/dev/null || true)
    [[ -n "$H_KEY_VAL" ]] && { echo "$H_KEY_VAL" > "${CONFIG_DIR}/host.key"; chmod 600 "${CONFIG_DIR}/host.key"; }

    HCRED_RESP=$(curl -sf -X POST "http://127.0.0.1:${CP_HTTP_PORT}/api/projects/${PROJECT_ID}/agent-grpc-credentials" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer ${JWT}" \
      -d '{"name":"host-agent-1"}' 2>&1) || { warn "Host credential failed — provision via dashboard later"; HCRED_RESP=""; }
    if [[ -n "$HCRED_RESP" ]]; then
      HC=$(echo "$HCRED_RESP"    | python3 -c "import sys,json;print(json.load(sys.stdin).get('ca_cert',''))"     2>/dev/null || true)
      HCERT=$(echo "$HCRED_RESP" | python3 -c "import sys,json;print(json.load(sys.stdin).get('client_cert',''))" 2>/dev/null || true)
      HKEY=$(echo "$HCRED_RESP"  | python3 -c "import sys,json;print(json.load(sys.stdin).get('client_key',''))"  2>/dev/null || true)
      if [[ -n "$HC" ]] && [[ -n "$HKEY" ]]; then
        echo "$HC"    > "${CERTS_DIR}/host/ca.crt"
        echo "$HCERT" > "${CERTS_DIR}/host/client.crt"
        echo "$HKEY"  > "${CERTS_DIR}/host/client.key"
        chmod 600 "${CERTS_DIR}/host/client.key"
        ok "Host agent certificates saved to ${CERTS_DIR}/host/"
      fi
    fi
  fi

  # Stop the temporary setup server
  python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGTERM)" 2>/dev/null || true
  sleep 2
  python3 -c "import os,signal; os.kill(${SETUP_PID}, signal.SIGKILL)" 2>/dev/null || true
fi

# ─── Agent-only config (no control plane on same box) ────────────────────────
if $INSTALL_WAF && ! $INSTALL_CP; then
  step "WAF Agent Configuration"
  echo ""
  echo "  Enter details from your WAFio Control Plane dashboard (Credentials page)."
  echo ""
  prompt "Control plane gRPC address (host:9090): "; read -r CP_GRPC_ADDR_WAF
  prompt "Agent name [waf-agent-1]: ";               read -r WAF_AGENT_NAME
  WAF_AGENT_NAME="${WAF_AGENT_NAME:-waf-agent-1}"
  echo ""
  echo "  Paste the CA certificate (press Enter twice when done):"
  WAF_CA_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; WAF_CA_TXT+="$LINE"$'\n'; done
  echo "  Paste the client certificate:"
  WAF_CERT_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; WAF_CERT_TXT+="$LINE"$'\n'; done
  echo "  Paste the client private key:"
  WAF_KEY_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; WAF_KEY_TXT+="$LINE"$'\n'; done
  printf '%s' "$WAF_CA_TXT"   > "${CERTS_DIR}/waf/ca.crt"
  printf '%s' "$WAF_CERT_TXT" > "${CERTS_DIR}/waf/client.crt"
  printf '%s' "$WAF_KEY_TXT"  > "${CERTS_DIR}/waf/client.key"
  chmod 600 "${CERTS_DIR}/waf/client.key"
  cat > "${CONFIG_DIR}/waf-agent.env" <<ENV_EOF
WAFIO_GRPC_SERVER=${CP_GRPC_ADDR_WAF}
WAFIO_AGENT_NAME=${WAF_AGENT_NAME}
WAFIO_GRPC_CA_CERT=${CERTS_DIR}/waf/ca.crt
WAFIO_GRPC_CLIENT_CERT=${CERTS_DIR}/waf/client.crt
WAFIO_GRPC_CLIENT_KEY=${CERTS_DIR}/waf/client.key
WAFIO_LICENSE_FILE=${CONFIG_DIR}/waf-agent/license.json
WAFIO_CRS_DEFAULT_DIR=${SECURITY_DATA_DIR}/crs/coreruleset/rules
WAFIO_CRS_SETUP_FILE=${SECURITY_DATA_DIR}/crs/coreruleset/crs-setup.conf
WAFIO_CRS_RULES_GLOB=${SECURITY_DATA_DIR}/crs/coreruleset/rules/*.conf
WAFIO_GEO_COUNTRY_DB_PATH=${SECURITY_DATA_DIR}/geoip/country.mmdb
WAFIO_GEO_CITY_DB_PATH=${SECURITY_DATA_DIR}/geoip/city.mmdb
WAFIO_GEO_ASN_DB_PATH=${SECURITY_DATA_DIR}/geoip/asn.mmdb
WAFIO_UA_BAD_BOT_FILE=${SECURITY_DATA_DIR}/user-agents/bot.txt
WAFIO_JA3_BLACKLIST_FILE=${SECURITY_DATA_DIR}/ja3/ja3.txt
ENV_EOF
  chmod 600 "${CONFIG_DIR}/waf-agent.env"
  ok "WAF agent environment file written"
fi

if $INSTALL_CP && $INSTALL_WAF; then
  cat > "${CONFIG_DIR}/waf-agent.env" <<ENV_EOF
WAFIO_GRPC_SERVER=127.0.0.1:${CP_GRPC_PORT}
WAFIO_GRPC_SERVER_NAME=${GRPC_DOMAIN}
WAFIO_AGENT_NAME=waf-agent-1
WAFIO_GRPC_CA_CERT=${CERTS_DIR}/waf/ca.crt
WAFIO_GRPC_CLIENT_CERT=${CERTS_DIR}/waf/client.crt
WAFIO_GRPC_CLIENT_KEY=${CERTS_DIR}/waf/client.key
WAFIO_LICENSE_FILE=${CONFIG_DIR}/waf-agent/license.json
WAFIO_CRS_DEFAULT_DIR=${SECURITY_DATA_DIR}/crs/coreruleset/rules
WAFIO_CRS_SETUP_FILE=${SECURITY_DATA_DIR}/crs/coreruleset/crs-setup.conf
WAFIO_CRS_RULES_GLOB=${SECURITY_DATA_DIR}/crs/coreruleset/rules/*.conf
WAFIO_GEO_COUNTRY_DB_PATH=${SECURITY_DATA_DIR}/geoip/country.mmdb
WAFIO_GEO_CITY_DB_PATH=${SECURITY_DATA_DIR}/geoip/city.mmdb
WAFIO_GEO_ASN_DB_PATH=${SECURITY_DATA_DIR}/geoip/asn.mmdb
WAFIO_UA_BAD_BOT_FILE=${SECURITY_DATA_DIR}/user-agents/bot.txt
WAFIO_JA3_BLACKLIST_FILE=${SECURITY_DATA_DIR}/ja3/ja3.txt
ENV_EOF
  chmod 600 "${CONFIG_DIR}/waf-agent.env"
fi

if $INSTALL_HOST && ! $INSTALL_CP; then
  step "Host Agent Configuration"
  echo ""
  echo "  Enter details from your WAFio Control Plane dashboard (Hosts + Credentials pages)."
  echo ""
  prompt "Control plane gRPC address (host:9090): "; read -r CP_GRPC_ADDR_HOST
  prompt "Host key (from dashboard → Hosts): ";      read -r H_KEY_INPUT
  echo ""
  echo "  Paste the CA certificate (press Enter twice when done):"
  HC_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; HC_TXT+="$LINE"$'\n'; done
  echo "  Paste the client certificate:"
  HCERT_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; HCERT_TXT+="$LINE"$'\n'; done
  echo "  Paste the client private key:"
  HKEY_TXT=""; while IFS= read -r LINE; do [[ -z "$LINE" ]] && break; HKEY_TXT+="$LINE"$'\n'; done
  printf '%s' "$HC_TXT"    > "${CERTS_DIR}/host/ca.crt"
  printf '%s' "$HCERT_TXT" > "${CERTS_DIR}/host/client.crt"
  printf '%s' "$HKEY_TXT"  > "${CERTS_DIR}/host/client.key"
  chmod 600 "${CERTS_DIR}/host/client.key"
  cat > "${CONFIG_DIR}/host-agent.env" <<ENV_EOF
WAFIO_GRPC_SERVER=${CP_GRPC_ADDR_HOST}
WAFIO_GRPC_CA_CERT=${CERTS_DIR}/host/ca.crt
WAFIO_GRPC_CLIENT_CERT=${CERTS_DIR}/host/client.crt
WAFIO_GRPC_CLIENT_KEY=${CERTS_DIR}/host/client.key
WAFIO_HOST_KEY=${H_KEY_INPUT}
WAFIO_LICENSE_FILE=${CONFIG_DIR}/host-agent/license.json
WAFIO_SERVER_URL=http://127.0.0.1:${CP_HTTP_PORT}
ENV_EOF
  chmod 600 "${CONFIG_DIR}/host-agent.env"
  ok "Host agent environment file written"
fi

if $INSTALL_CP && $INSTALL_HOST; then
  H_KEY_VAL=$(cat "${CONFIG_DIR}/host.key" 2>/dev/null || echo "")
  cat > "${CONFIG_DIR}/host-agent.env" <<ENV_EOF
WAFIO_GRPC_SERVER=127.0.0.1:${CP_GRPC_PORT}
WAFIO_GRPC_SERVER_NAME=${GRPC_DOMAIN}
WAFIO_GRPC_CA_CERT=${CERTS_DIR}/host/ca.crt
WAFIO_GRPC_CLIENT_CERT=${CERTS_DIR}/host/client.crt
WAFIO_GRPC_CLIENT_KEY=${CERTS_DIR}/host/client.key
WAFIO_HOST_KEY=${H_KEY_VAL}
WAFIO_LICENSE_FILE=${CONFIG_DIR}/host-agent/license.json
WAFIO_SERVER_URL=http://127.0.0.1:${CP_HTTP_PORT}
ENV_EOF
  chmod 600 "${CONFIG_DIR}/host-agent.env"
fi

# ─── 12. Systemd services ─────────────────────────────────────────────────────
step "Installing Systemd Services"

if $INSTALL_CP; then
  cat > /etc/systemd/system/wafio.service <<SVC_EOF
[Unit]
Description=WAFio Control Plane
After=network.target postgresql.service

[Service]
Type=simple
User=root
WorkingDirectory=${INSTALL_DIR}
ExecStart=/usr/local/bin/wafio serve
Restart=on-failure
RestartSec=5s
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wafio
LimitNOFILE=65536

[Install]
WantedBy=multi-user.target
SVC_EOF
  systemctl daemon-reload
  systemctl enable wafio
  systemctl start wafio
  sleep 3
  systemctl is-active --quiet wafio \
    && ok "wafio control plane service started" \
    || warn "wafio service issue — check: journalctl -u wafio -n 50"
fi

if $INSTALL_WAF && [[ -s "${CERTS_DIR}/waf/client.key" ]]; then
  cat > /etc/systemd/system/wafio-waf-agent.service <<SVC_EOF
[Unit]
Description=WAFio WAF Agent
After=network.target

[Service]
Type=simple
User=root
EnvironmentFile=${CONFIG_DIR}/waf-agent.env
ExecStart=/usr/local/bin/wafio-waf-agent
Restart=on-failure
RestartSec=5s
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wafio-waf-agent
LimitNOFILE=65536

[Install]
WantedBy=multi-user.target
SVC_EOF
  systemctl daemon-reload
  systemctl enable wafio-waf-agent
  systemctl start wafio-waf-agent
  sleep 2
  systemctl is-active --quiet wafio-waf-agent \
    && ok "wafio-waf-agent service started" \
    || warn "wafio-waf-agent issue — check: journalctl -u wafio-waf-agent -n 30"
elif $INSTALL_WAF; then
  warn "WAF agent certificates not found — provision credentials from dashboard, then:"
  warn "  sudo systemctl start wafio-waf-agent"
fi

if $INSTALL_HOST && [[ -s "${CERTS_DIR}/host/client.key" ]]; then
  H_KEY_VAL=$(cat "${CONFIG_DIR}/host.key" 2>/dev/null || echo "")
  _NIC=$(ip -o link show | grep -v 'lo:' | grep 'state UP' | head -1 | awk -F': ' '{print $2}' | awk '{print $1}' 2>/dev/null || echo "eth0")
  _EBPF_DIR="${INSTALL_DIR}/ebpf"
  cat > /etc/systemd/system/wafio-host-agent.service <<SVC_EOF
[Unit]
Description=WAFio Host Agent (eBPF)
After=network.target

[Service]
Type=simple
User=root
EnvironmentFile=${CONFIG_DIR}/host-agent.env
ExecStart=/usr/local/bin/wafio-host-agent \
  --interface=${_NIC} \
  --geoip-db=${SECURITY_DATA_DIR}/geoip/country.mmdb \
  --xdp-obj=${_EBPF_DIR}/wafio_xdp.o \
  --tc-obj=${_EBPF_DIR}/wafio_tc.o \
  --trace-obj=${_EBPF_DIR}/wafio_trace.o \
  --ebpf
Restart=on-failure
RestartSec=5s
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wafio-host-agent
LimitNOFILE=65536

[Install]
WantedBy=multi-user.target
SVC_EOF
  systemctl daemon-reload
  systemctl enable wafio-host-agent
  systemctl start wafio-host-agent
  sleep 3
  systemctl is-active --quiet wafio-host-agent \
    && ok "wafio-host-agent service started" \
    || warn "wafio-host-agent issue — check: journalctl -u wafio-host-agent -n 30"
elif $INSTALL_HOST; then
  warn "Host agent certificates not found — provision credentials from dashboard, then:"
  warn "  sudo systemctl start wafio-host-agent"
fi

# ─── 13. Health check ─────────────────────────────────────────────────────────
if $INSTALL_CP; then
  step "Final Health Check"
  HOST_IP=$(hostname -I | awk '{print $1}')
  HEALTHY=false
  for i in $(seq 1 20); do
    curl -sf "http://127.0.0.1:${CP_HTTP_PORT}/health" &>/dev/null && { HEALTHY=true; break; }
    sleep 2
  done
  $HEALTHY \
    && ok "Control plane is healthy at http://${HOST_IP}:${CP_HTTP_PORT}" \
    || warn "Health check failed — check: journalctl -u wafio -n 50"
fi

# ─── Done ─────────────────────────────────────────────────────────────────────
HOST_IP=$(hostname -I | awk '{print $1}' 2>/dev/null || echo "YOUR_SERVER_IP")

echo ""
hr
echo -e "  ${GREEN}${BOLD}✓  WAFio Installation Complete!${NC}"
hr
echo ""
$INSTALL_CP && echo -e "  ${BOLD}Dashboard URL:${NC}    http://${HOST_IP}:${CP_HTTP_PORT}"
$INSTALL_CP && [[ -n "$ADMIN_EMAIL" ]]  && echo -e "  ${BOLD}Admin email:${NC}      ${ADMIN_EMAIL}"
$INSTALL_CP && [[ -n "$PROJECT_ID" ]]   && echo -e "  ${BOLD}Project ID:${NC}       ${PROJECT_ID}"
echo ""
echo -e "  ${BOLD}Service status:${NC}"
$INSTALL_CP   && echo "    sudo systemctl status wafio"
$INSTALL_WAF  && echo "    sudo systemctl status wafio-waf-agent"
$INSTALL_HOST && echo "    sudo systemctl status wafio-host-agent"
echo ""
echo -e "  ${BOLD}View logs:${NC}"
$INSTALL_CP   && echo "    sudo journalctl -u wafio -f"
$INSTALL_WAF  && echo "    sudo journalctl -u wafio-waf-agent -f"
$INSTALL_HOST && echo "    sudo journalctl -u wafio-host-agent -f"
echo ""
echo -e "  ${BOLD}Documentation:${NC}   https://wafio.cloud/docs"
echo ""
hr
echo ""
