#!/usr/bin/env bash
set -euo pipefail

# Rizoma AV Server Setup
#
# Enrolls a dedicated scanner host into the Rizoma mesh via the
# dashboard's Add AV Server flow. Clone of relay-setup.sh mechanics
# (lifecycle, repo bootstrap, /dev/tty auth-key prompt, exchange):
#
#   1. Signals /v1/install/start so the dashboard's modal moves
#      out of "waiting".
#   2. Installs rizoma-av from the package repo (apt or dnf,
#      detected automatically).
#   3. Prompts for the one-time auth_key (60s TTL, single-use,
#      purpose add_av_server, minted by the dashboard) — never
#      passed on the command line.
#   4. Exchanges the auth_key at /v1/avserver/enroll-exchange —
#      the coordinator creates the av-server row and returns
#      server_id + server_secret + mesh cert material.
#   5. Writes /etc/rizoma/av.env + state and starts
#      rizoma-av.service. From here the server heartbeats
#      presence/load and serves scan verdicts to agents.
#
# AV servers take no inbound internet traffic (mesh dials only),
# so public_addr is the LAN-reachable host:port agents dial
# (default-route source IP auto-detected, port 8443 default).
#
# Usage (the command the Add AV Server modal renders):
#   curl -fsSL https://COORDINATOR/install/av-setup.sh | bash -s -- \
#     --coordinator-url 'https://COORDINATOR' \
#     --session-id 'SESSION' \
#     --name 'av-lan-1' \
#     --public-addr '192.168.50.111:8443'  (optional — auto-detected) \
#     --region 'lan'

AV_NAME=""
PUBLIC_ADDR=""
REGION="default"
COORDINATOR_URL=""
INSTALL_SESSION_ID=""
AUTH_KEY=""
LISTEN_PORT="8443"
MAX_SCANS=8
REPO_URL="${REPO_URL:-https://repo.rizomarl.com/mesh}"
CHANNEL="${CHANNEL:-stable}"
ENV_PATH="/etc/rizoma/av.env"
LOG_FILE="/var/log/rizoma-av-setup.log"

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$LOG_FILE" >&2; }
die() { log "ERROR: $*"; exit 1; }
json_escape() { printf '%s' "$1" | sed -e 's/\\/\\\\/g' -e 's/"/\\"/g'; }

mkdir -p "$(dirname "$LOG_FILE")" 2>/dev/null || LOG_FILE="/tmp/rizoma-av-setup.log"
touch "$LOG_FILE" 2>/dev/null || true

if [[ "$(id -u)" -ne 0 ]]; then
  echo "error: av setup must run as root (sudo)" >&2
  exit 1
fi

while [[ $# -gt 0 ]]; do
  case "$1" in
    --coordinator-url) COORDINATOR_URL="${2:-}"; shift 2 ;;
    --session-id) INSTALL_SESSION_ID="${2:-}"; shift 2 ;;
    --name) AV_NAME="${2:-}"; shift 2 ;;
    --public-addr) PUBLIC_ADDR="${2:-}"; shift 2 ;;
    --region) REGION="${2:-}"; shift 2 ;;
    --auth-key) AUTH_KEY="${2:-}"; shift 2 ;;
    --listen-port) LISTEN_PORT="${2:-8443}"; shift 2 ;;
    --max-scans) MAX_SCANS="${2:-8}"; shift 2 ;;
    *) echo "unknown argument: $1" >&2; shift ;;
  esac
done

if [[ -z "$COORDINATOR_URL" ]]; then
  read -r -p "Coordinator URL (https://mesh.example.com): " COORDINATOR_URL
fi
COORDINATOR_URL="${COORDINATOR_URL%/}"
[[ -n "$COORDINATOR_URL" ]] || die "coordinator URL is required"
[[ "$COORDINATOR_URL" == http* ]] || die "coordinator URL must start with http(s)://"

if [[ -z "$AV_NAME" ]]; then
  read -r -p "AV server name (av-lan-1): " AV_NAME
fi
[[ -n "$AV_NAME" ]] || die "server name is required"

# --- Reachable address: AV servers are dialed over the mesh/LAN,
# not the internet, so default-route source IP wins over echo
# services (reverse priority vs relay-setup.sh).
detect_public_addr() {
  local ip
  ip="$(ip route get 1.1.1.1 2>/dev/null | sed -n 's/.*src \([0-9.]*\).*/\1/p' | head -n1)"
  if [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    echo "$ip"
    return 0
  fi
  for svc in "https://api.ipify.org" "https://ifconfig.me/ip"; do
    ip="$(curl -fsS --max-time 5 "$svc" 2>/dev/null | tr -d '[:space:]')"
    if [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
      echo "$ip"
      return 0
    fi
  done
  return 1
}

if [[ -z "$PUBLIC_ADDR" ]]; then
  if detected="$(detect_public_addr)"; then
    PUBLIC_ADDR="${detected}:${LISTEN_PORT}"
    log "Auto-detected reachable address: $PUBLIC_ADDR"
  else
    read -r -p "Reachable address (host:port, e.g. 192.168.50.111:8443): " PUBLIC_ADDR
  fi
fi
[[ "$PUBLIC_ADDR" == *:* ]] || die "public address must be host:port"

install_lifecycle_call() {
  local endpoint="$1"
  local body="$2"
  [[ -z "$INSTALL_SESSION_ID" ]] && return 0
  command -v curl >/dev/null 2>&1 || return 0
  local url="${COORDINATOR_URL%/}/v1/install/${endpoint}"
  curl -sS --max-time 10 \
    -H "Content-Type: application/json" \
    -X POST \
    --data "$body" \
    "$url" >/dev/null 2>&1 || true
}

install_lifecycle_start() {
  install_lifecycle_call "start" "{\"session_id\":\"$(json_escape "$INSTALL_SESSION_ID")\",\"hostname\":\"$(json_escape "$(hostname)")\",\"created_via\":\"av_setup\"}"
}

install_lifecycle_ready() {
  install_lifecycle_call "ready" "{\"session_id\":\"$(json_escape "$INSTALL_SESSION_ID")\"}"
}

install_lifecycle_complete() {
  local node_id="$1"
  install_lifecycle_call "complete" "{\"session_id\":\"$(json_escape "$INSTALL_SESSION_ID")\",\"node_id\":\"$(json_escape "$node_id")\"}"
}

install_lifecycle_failed() {
  local reason="$1"
  install_lifecycle_call "failed" "{\"session_id\":\"$(json_escape "$INSTALL_SESSION_ID")\",\"reason\":\"$(json_escape "$reason")\"}"
}

install_lifecycle_start

# --- Package installation (jq hard dependency for parsing) ---
repo_already_configured() {
  if command -v apt-get >/dev/null 2>&1; then
    [[ -f /etc/apt/sources.list.d/rizoma.list ]]
  elif command -v dnf >/dev/null 2>&1; then
    [[ -f /etc/yum.repos.d/rizoma.repo ]]
  else
    false
  fi
}

install_apt_keyring() {
  local tmp_key tmp_out
  tmp_key="$(mktemp /tmp/rizoma-keyring.XXXXXX)"
  if ! curl -fsSL --max-time 15 "$REPO_URL/keys/rizoma-archive-keyring.gpg" -o "$tmp_key"; then
    rm -f "$tmp_key"
    die "could not fetch the rizoma repo keyring from $REPO_URL"
  fi
  if grep -q "BEGIN PGP PUBLIC KEY BLOCK" "$tmp_key"; then
    tmp_out="$(mktemp /tmp/rizoma-keyring-out.XXXXXX)"
    if ! gpg --batch --yes --dearmor -o "$tmp_out" "$tmp_key" 2>/dev/null; then
      rm -f "$tmp_key" "$tmp_out"
      die "gpg dearmor of the rizoma keyring failed"
    fi
    install -m 0644 "$tmp_out" /usr/share/keyrings/rizoma-archive-keyring.gpg
    rm -f "$tmp_out"
  else
    install -m 0644 "$tmp_key" /usr/share/keyrings/rizoma-archive-keyring.gpg
  fi
  rm -f "$tmp_key"
}

bootstrap_repo() {
  if command -v apt-get >/dev/null 2>&1; then
    mkdir -p /usr/share/keyrings
    install_apt_keyring
    printf 'deb [signed-by=/usr/share/keyrings/rizoma-archive-keyring.gpg] %s/apt/%s ./\n' "$REPO_URL" "$CHANNEL" > /etc/apt/sources.list.d/rizoma.list
  elif command -v dnf >/dev/null 2>&1; then
    mkdir -p /etc/pki/rpm-gpg /etc/yum.repos.d
    local tmp_key
    tmp_key="$(mktemp /tmp/rizoma-rpm-key.XXXXXX)"
    if ! curl -fsSL --max-time 15 "$REPO_URL/keys/RPM-GPG-KEY-rizoma" -o "$tmp_key"; then
      rm -f "$tmp_key"
      die "could not fetch the rizoma RPM key from $REPO_URL"
    fi
    install -m 0644 "$tmp_key" /etc/pki/rpm-gpg/RPM-GPG-KEY-rizoma
    rm -f "$tmp_key"
    printf '[rizoma]\nname=Rizoma\nbaseurl=%s/dnf/%s/\nenabled=1\ngpgcheck=1\ngpgkey=file:///etc/pki/rpm-gpg/RPM-GPG-KEY-rizoma\n' "$REPO_URL" "$CHANNEL" > /etc/yum.repos.d/rizoma.repo
  else
    die "neither apt-get nor dnf found on this host"
  fi
}

install_package() {
  if command -v apt-get >/dev/null 2>&1; then
    export DEBIAN_FRONTEND=noninteractive
    if ! repo_already_configured; then
      log "Configuring rizoma package repo (first run on this host)..."
      bootstrap_repo
    fi
    if ! apt-get update -qq; then
      log "apt-get update reported issues (continuing — repo metadata may be partially cached)"
    fi
    if ! apt-get install -y rizoma-av jq; then
      die "apt install rizoma-av jq failed — see the apt output above"
    fi
  elif command -v dnf >/dev/null 2>&1; then
    if ! repo_already_configured; then
      log "Configuring rizoma package repo (first run on this host)..."
      bootstrap_repo
    fi
    if ! dnf install -y rizoma-av jq; then
      die "dnf install rizoma-av jq failed — see the dnf output above"
    fi
  else
    die "neither apt-get nor dnf found on this host"
  fi
}

if ! id -u rizoma-av >/dev/null 2>&1; then
  log "Creating rizoma-av system user..."
  useradd --system --no-create-home --shell /usr/sbin/nologin rizoma-av || die "failed to create rizoma-av system user"
fi

if ! command -v rizoma-av >/dev/null 2>&1 || ! command -v jq >/dev/null 2>&1; then
  log "Installing rizoma-av + jq packages..."
  install_package
fi
command -v rizoma-av >/dev/null 2>&1 || die "rizoma-av binary not found after install"
command -v jq >/dev/null 2>&1 || die "jq not found after install"

# --- Auth key prompt AFTER install (dashboard in "ready" state) ---
install_lifecycle_ready

prompt_for_auth_key() {
  if [[ -n "$AUTH_KEY" ]]; then
    return 0
  fi
  local tty_in="/dev/stdin"
  if [[ ! -t 0 ]] && [[ -r /dev/tty ]]; then
    tty_in="/dev/tty"
  fi
  local attempt=0
  while [[ $attempt -lt 3 ]]; do
    attempt=$((attempt + 1))
    printf 'Paste the one-time auth key from the dashboard Add AV Server modal: ' > "$tty_in" 2>/dev/null || printf 'Paste the one-time auth key: '
    IFS= read -r -t 180 AUTH_KEY < "$tty_in" || {
      echo "" >&2
      die "no auth key provided (timed out or unreadable input)"
    }
    AUTH_KEY="$(printf '%s' "$AUTH_KEY" | tr -d '[:space:]')"
    if [[ "$AUTH_KEY" =~ ^ak_[A-Z0-9]{16}$ ]]; then
      return 0
    fi
    log "auth key format invalid (expected ak_ + 16 chars), try again ($attempt/3)"
    AUTH_KEY=""
  done
  die "auth key required — mint a fresh one from the dashboard (60s TTL)"
}

prompt_for_auth_key

# --- Enrollment exchange ---
EXCHANGE_BODY="$(jq -nc \
  --arg auth_key "$AUTH_KEY" \
  --arg name "$AV_NAME" \
  --arg public_addr "$PUBLIC_ADDR" \
  --arg region "$REGION" \
  --argjson max_scans "$MAX_SCANS" \
  '{auth_key:$auth_key,name:$name,public_addr:$public_addr,region:$region,max_scans:$max_scans}')"
EXCHANGE_RESP="$(curl -fsS --max-time 30 \
  -H "Content-Type: application/json" \
  -X POST \
  --data "$EXCHANGE_BODY" \
  "${COORDINATOR_URL%/}/v1/avserver/enroll-exchange")" || {
  install_lifecycle_failed "enroll-exchange request failed"
  die "enroll exchange failed — is the auth key fresh (60s TTL) and purpose add_av_server?"
}

SERVER_ID="$(printf '%s' "$EXCHANGE_RESP" | jq -r '.server_id // empty')"
SERVER_SECRET="$(printf '%s' "$EXCHANGE_RESP" | jq -r '.server_secret // empty')"
[[ -n "$SERVER_ID" && -n "$SERVER_SECRET" ]] || {
  install_lifecycle_failed "enroll-exchange response missing identity"
  die "enroll exchange response missing server identity: $EXCHANGE_RESP"
}

# --- Persist env + identity material ---
install -d -m 0711 /etc/rizoma
install -d -m 0700 /var/lib/rizoma/av
CERT_DIR="/var/lib/rizoma/av/certs"
install -d -m 0755 "$CERT_DIR"
printf '%s' "$EXCHANGE_RESP" | jq -r '.cert // empty' > "$CERT_DIR/av_cert.pem"
printf '%s' "$EXCHANGE_RESP" | jq -r '.key // empty' > "$CERT_DIR/av_key.pem"
{
  printf '%s' "$EXCHANGE_RESP" | jq -r '.ca_cert // empty'
  printf '%s' "$EXCHANGE_RESP" | jq -r '.ca_cert_prev // empty'
} > "$CERT_DIR/mesh_ca.pem"
chmod 0600 "$CERT_DIR/av_cert.pem" "$CERT_DIR/av_key.pem" "$CERT_DIR/mesh_ca.pem"

cat > "$ENV_PATH" <<EOF
RIZOMA_AV_COORDINATOR_URL=$COORDINATOR_URL
RIZOMA_AV_SERVER_ID=$SERVER_ID
RIZOMA_AV_SERVER_SECRET=$SERVER_SECRET
RIZOMA_AV_PUBLIC_ADDR=$PUBLIC_ADDR
RIZOMA_AV_REGION=$REGION
RIZOMA_AV_MAX_SCANS=$MAX_SCANS
EOF
chmod 0600 "$ENV_PATH"

# --- First heartbeat before declaring success (fail fast on bad creds) ---
LISTEN_PORT="${PUBLIC_ADDR##*:}"
HB_BODY="$(jq -nc --arg id "$SERVER_ID" '{server_id:$id,active_scans:0,max_scans:8}')"
if ! curl -fsS --max-time 15 \
  -H "Content-Type: application/json" \
  -H "X-Rizoma-Av-Secret: $SERVER_SECRET" \
  -X POST \
  --data "$HB_BODY" \
  "${COORDINATOR_URL%/}/v1/avserver/heartbeat" >/dev/null; then
  install_lifecycle_failed "first heartbeat failed"
  die "first heartbeat failed — enrollment written but server unreachable; check coordinator URL"
fi

systemctl daemon-reload >/dev/null 2>&1 || true
systemctl enable --now rizoma-av.service || die "failed to start rizoma-av.service"
install_lifecycle_complete "$SERVER_ID"
log "AV server enrolled as $SERVER_ID ($PUBLIC_ADDR). Service running."