#!/usr/bin/env bash
# Vetting quick-register.
#
# Run on the target host (any Linux with root) before the host is wiped:
#   curl -fsSL {{.OrchestratorURL}}/register/quick.sh | sudo bash
#
# Detects the primary NIC's MAC, probes hardware (CPU / RAM / disks /
# NICs / GPUs) into an expected-spec YAML, and POSTs everything to
# {{.OrchestratorURL}}/api/v1/hosts. After registration, go to the
# orchestrator's dashboard and click "Start vetting" for the new host.
#
# Env overrides (all optional):
#   NAME             Host display name (default: `hostname -s`)
#   MAC              Force a specific MAC (default: autodetect)
#   WOL_BROADCAST    WoL broadcast IP (default: primary iface broadcast)
#   WOL_PORT         WoL UDP port (default: 9)
#   NOTES            Free-text notes
#   ORCH_URL         Override orchestrator base URL
#   INSTALL_AGENT    1=install vetting-reporter systemd service (default)
#                    0=skip the agent install (registration only)
#                    Pass via:  curl ... | sudo INSTALL_AGENT=0 bash
set -euo pipefail

ORCH_URL="${ORCH_URL:-{{.OrchestratorURL}}}"

if [[ -z "${ORCH_URL}" ]]; then
    echo "ERROR: ORCH_URL is empty; pass it via env." >&2
    exit 1
fi

primary_iface() {
    # Prefer the interface carrying the default route — that's the
    # canonical "primary" iface (e.g. vmbr0 on Proxmox, eno1 on bare
    # metal). `ip link show` order picks the physical NIC on Proxmox,
    # but that NIC has no IPv4, so we'd miss the broadcast address.
    local iface
    iface="$(ip -o -4 route show default 2>/dev/null \
        | awk '{for(i=1;i<=NF;i++) if($i=="dev") {print $(i+1); exit}}')"
    if [[ -n "${iface}" ]]; then
        echo "${iface}"
        return
    fi
    # Fallback: first non-virtual interface that has an IPv4 address.
    ip -o -4 addr show 2>/dev/null \
        | awk '{
                 name=$2
                 if (name ~ /^(lo|docker|br-|veth|virbr|bond|tun|tap|fwbr|fwpr|fwln|wlan|wlp)/) next
                 print name; exit
               }'
}

# compute_broadcast "192.168.1.250/24" → "192.168.1.255"
compute_broadcast() {
    local cidr="$1" ip prefix a b c d host mask inv bc
    ip="${cidr%/*}"
    prefix="${cidr#*/}"
    [[ "${ip}" == *.*.*.* && "${prefix}" =~ ^[0-9]+$ ]] || return 1
    IFS=. read -r a b c d <<<"${ip}"
    host=$(( (a<<24) | (b<<16) | (c<<8) | d ))
    mask=$(( (0xFFFFFFFF << (32 - prefix)) & 0xFFFFFFFF ))
    inv=$(( (~mask) & 0xFFFFFFFF ))
    bc=$(( host | inv ))
    printf '%d.%d.%d.%d' $(( (bc>>24)&0xFF )) $(( (bc>>16)&0xFF )) $(( (bc>>8)&0xFF )) $(( bc&0xFF ))
}

IFACE="$(primary_iface || true)"
if [[ -z "${IFACE}" ]]; then
    echo "ERROR: could not pick a primary network interface." >&2
    exit 1
fi

NAME="${NAME:-$(hostname -s 2>/dev/null || hostname)}"
MAC="${MAC:-$(cat /sys/class/net/${IFACE}/address 2>/dev/null || true)}"
if [[ -z "${MAC}" ]]; then
    echo "ERROR: could not read MAC for ${IFACE}." >&2
    exit 1
fi

if [[ -z "${WOL_BROADCAST:-}" ]]; then
    ipinfo="$(ip -o -4 addr show dev "${IFACE}" 2>/dev/null || true)"
    WOL_BROADCAST="$(awk '{for(i=1;i<=NF;i++) if($i=="brd") {print $(i+1); exit}}' <<<"${ipinfo}")"
    # Bridges (vmbr0 on Proxmox, br0 on Linux bridges) often report
    # brd 0.0.0.0 or omit it entirely. Compute from the inet CIDR
    # before giving up on 255.255.255.255.
    if [[ -z "${WOL_BROADCAST}" || "${WOL_BROADCAST}" == "0.0.0.0" ]]; then
        cidr="$(awk '{for(i=1;i<=NF;i++) if($i=="inet") {print $(i+1); exit}}' <<<"${ipinfo}")"
        if [[ "${cidr}" == */* ]]; then
            WOL_BROADCAST="$(compute_broadcast "${cidr}" || true)"
        fi
    fi
fi
WOL_BROADCAST="${WOL_BROADCAST:-255.255.255.255}"
WOL_PORT="${WOL_PORT:-9}"

# --- Hardware probes ---
cores="$(nproc 2>/dev/null || echo 0)"
cpu_model="$(awk -F: '/^model name/ {sub(/^ */, "", $2); print $2; exit}' /proc/cpuinfo 2>/dev/null || true)"
mem_gib="$(awk '/^MemTotal:/ {printf "%d", ($2/1024/1024) + 0.5; exit}' /proc/meminfo 2>/dev/null || echo 0)"

disk_yaml=""
while read -r name size serial; do
    [[ -z "${name}" ]] && continue
    [[ "${name}" =~ ^(sd|nvme|vd|hd) ]] || continue
    [[ -z "${serial}" ]] && continue
    size_gb=$(( size / 1000 / 1000 / 1000 ))
    disk_yaml+="  - serial: \"${serial}\"
    size_gb: ${size_gb}
"
done < <(lsblk -dn -b -o NAME,SIZE,SERIAL 2>/dev/null || true)

nic_yaml=""
for iface_dir in /sys/class/net/*; do
    iface_name="$(basename "${iface_dir}")"
    [[ "${iface_name}" =~ ^(lo|docker|br-|veth|virbr|bond|tun|tap)$ ]] && continue
    nic_mac="$(cat "${iface_dir}/address" 2>/dev/null || true)"
    [[ -z "${nic_mac}" || "${nic_mac}" == "00:00:00:00:00:00" ]] && continue
    speed_mbps="$(cat "${iface_dir}/speed" 2>/dev/null || echo 0)"
    if [[ "${speed_mbps}" =~ ^[0-9]+$ ]] && (( speed_mbps > 0 )); then
        speed_gbps=$(( (speed_mbps + 999) / 1000 ))
    else
        speed_gbps=0
    fi
    nic_yaml+="  - mac: \"${nic_mac}\"
    speed_gbps: ${speed_gbps}
"
done

gpu_yaml=""
if command -v lspci >/dev/null 2>&1; then
    while IFS= read -r gpu; do
        [[ -z "${gpu}" ]] && continue
        gpu_yaml+="  - model: \"${gpu}\"
"
    done < <(lspci -mm 2>/dev/null | awk -F\" '/"(VGA|3D|Display)/ {print $6}')
fi

# Assemble the YAML spec. Empty sections are omitted so the diff engine
# skips them rather than demanding empty arrays on the actual side.
spec="cpu:
  model: \"${cpu_model}\"
  logical_cores: ${cores}
memory:
  total_gib: ${mem_gib}
"
[[ -n "${disk_yaml}" ]] && spec+="disks:
${disk_yaml}"
[[ -n "${nic_yaml}" ]] && spec+="nics:
${nic_yaml}"
[[ -n "${gpu_yaml}" ]] && spec+="gpus:
${gpu_yaml}"

# --- JSON escape (backslash, double-quote, newline, tab, CR) ---
json_escape() {
    local s="$1"
    s="${s//\\/\\\\}"
    s="${s//\"/\\\"}"
    s="${s//$'\t'/\\t}"
    s="${s//$'\r'/\\r}"
    s="${s//$'\n'/\\n}"
    printf '%s' "${s}"
}

payload=$(cat <<EOF
{
  "name": "$(json_escape "${NAME}")",
  "mac": "$(json_escape "${MAC}")",
  "wol_broadcast_ip": "$(json_escape "${WOL_BROADCAST}")",
  "wol_port": ${WOL_PORT},
  "expected_spec_yaml": "$(json_escape "${spec}")",
  "notes": "$(json_escape "${NOTES:-registered via quick-register on $(date -Is)}")"
}
EOF
)

echo "Registering ${NAME} (${MAC}) with ${ORCH_URL}..."
resp="$(curl -fsS -X POST \
    -H 'Content-Type: application/json' \
    -d "${payload}" \
    "${ORCH_URL}/api/v1/hosts")"
echo "OK: ${resp}"

# --- Optional: install the vetting-reporter systemd service so the
# host keeps heartbeating to the orchestrator long-term. Skipped when
# INSTALL_AGENT=0 or when systemctl isn't present (non-systemd hosts).
install_agent() {
    if [[ "${INSTALL_AGENT:-1}" == "0" ]]; then
        echo "Skipping agent install (INSTALL_AGENT=0)."
        return
    fi
    if ! command -v systemctl >/dev/null 2>&1; then
        echo "systemctl not found — skipping agent install."
        return
    fi
    echo "Installing vetting-reporter service..."
    install -d /etc/vetting /usr/local/bin
    # Download to a staging path, then use `install` — it unlinks the
    # target before writing, so re-running quick.sh on a host where
    # vetting-reporter is already running doesn't hit ETXTBSY
    # (curl error 23: write error on a busy executable).
    local staged="/usr/local/bin/.vetting-agent.new"
    if ! curl -fsSL "${ORCH_URL}/assets/vetting-agent-linux-amd64" -o "${staged}"; then
        echo "WARN: could not download agent from ${ORCH_URL}/assets/vetting-agent-linux-amd64"
        echo "WARN: registration succeeded but the host won't heartbeat."
        rm -f "${staged}"
        return
    fi
    install -m 0755 "${staged}" /usr/local/bin/vetting-agent
    rm -f "${staged}"
    cat >/etc/vetting/host-agent.yaml <<YAML
orchestrator_url: "${ORCH_URL}"
mac: "${MAC}"
interval: "30s"
YAML
    cat >/etc/systemd/system/vetting-reporter.service <<'UNIT'
[Unit]
Description=Vetting host-mode reporter
After=network-online.target
Wants=network-online.target

[Service]
ExecStart=/usr/local/bin/vetting-agent host -config /etc/vetting/host-agent.yaml
Restart=on-failure
RestartSec=5

[Install]
WantedBy=multi-user.target
UNIT
    systemctl daemon-reload
    systemctl enable vetting-reporter.service
    # `restart` instead of `start`: if the service was already running
    # from a prior quick.sh, this picks up the newly-installed binary.
    systemctl restart vetting-reporter.service
    echo "vetting-reporter.service enabled."
}
install_agent

echo
echo "Open ${ORCH_URL}/ and click 'Start vetting' on ${NAME}."