#!/bin/bash
# Nexora Cloud PC — node bootstrap / enrollment agent
# Usage:
#   CONTROL_PLANE_URL=https://cloud-gaming.nexorahost.com \
#   ENROLLMENT_CODE=your-code \
#   bash nexora-node-bootstrap.sh
set -euo pipefail

CONTROL_PLANE_URL="${CONTROL_PLANE_URL:-__CONTROL_PLANE_URL__}"
ENROLLMENT_CODE="${ENROLLMENT_CODE:-}"
STATE_DIR="${STATE_DIR:-/var/lib/nexora-node}"
TOKEN_FILE="${STATE_DIR}/node_token"
INTERVAL="${HEARTBEAT_INTERVAL:-30}"

if [[ "$CONTROL_PLANE_URL" == "__CONTROL_PLANE_URL__" || -z "$CONTROL_PLANE_URL" ]]; then
  CONTROL_PLANE_URL="https://cloud-gaming.nexorahost.com"
fi
CONTROL_PLANE_URL="${CONTROL_PLANE_URL%/}"

need() { command -v "$1" >/dev/null 2>&1 || { echo "Missing dependency: $1" >&2; exit 1; }; }
need curl
need hostname
need python3

mkdir -p "$STATE_DIR"

collect_checks() {
  python3 - <<'PY'
import json, os, platform, shutil, subprocess

def sh(cmd):
    try:
        return subprocess.check_output(cmd, shell=True, text=True, stderr=subprocess.DEVNULL).strip()
    except Exception:
        return ""

mem_kb = 0
try:
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemTotal:"):
                mem_kb = int(line.split()[1])
                break
except Exception:
    pass

cpu_cores = os.cpu_count() or 0
virt = sh("grep -E -m1 'vmx|svm' /proc/cpuinfo || true")
kvm_ok = os.path.exists("/dev/kvm")
gpu = sh("lspci 2>/dev/null | grep -iE 'vga|3d|display' | head -5")
disk = sh("df -BG --output=size,avail,target / | tail -1")

checks = {
    "hostname": platform.node(),
    "os": platform.platform(),
    "kernel": platform.release(),
    "arch": platform.machine(),
    "cpu_cores": cpu_cores,
    "memory_mb": mem_kb // 1024,
    "cpu_virt_flags": bool(virt),
    "kvm_device": kvm_ok,
    "libvirt_installed": shutil.which("virsh") is not None,
    "qemu_installed": shutil.which("qemu-system-x86_64") is not None,
    "gpu_lspci": [g.strip() for g in gpu.splitlines() if g.strip()],
    "root_disk": disk,
    "ready_for_kvm": bool(virt) or kvm_ok,
}
print(json.dumps(checks))
PY
}

api() {
  local method="$1" path="$2" data="${3:-}"
  local args=(-sS -X "$method" -H "Content-Type: application/json")
  if [[ -f "$TOKEN_FILE" ]]; then
    args+=(-H "Authorization: Bearer $(cat "$TOKEN_FILE")")
  fi
  if [[ -n "$data" ]]; then
    args+=(-d "$data")
  fi
  curl "${args[@]}" "${CONTROL_PLANE_URL}${path}"
}

enroll_if_needed() {
  if [[ -f "$TOKEN_FILE" ]]; then
    echo "Using existing node token at $TOKEN_FILE"
    return
  fi
  if [[ -z "$ENROLLMENT_CODE" ]]; then
    echo "Set ENROLLMENT_CODE to enroll this host." >&2
    exit 1
  fi
  local checks hostname payload resp
  checks="$(collect_checks)"
  hostname="$(hostname -f 2>/dev/null || hostname)"
  payload="$(python3 -c 'import json,sys; c=json.loads(sys.argv[1]); print(json.dumps({"enrollment_code":sys.argv[2],"hostname":sys.argv[3],"checks":c}))' "$checks" "$ENROLLMENT_CODE" "$hostname")"
  echo "Enrolling with control plane ${CONTROL_PLANE_URL} ..."
  resp="$(api POST /api/v1/nodes/enroll "$payload")"
  python3 -c 'import json,sys; d=json.loads(sys.argv[1]);
assert "node_token" in d, d
open(sys.argv[2],"w").write(d["node_token"]);
print("Enrolled node_id=%s status=%s desired=%s" % (d.get("node_id"), d.get("status"), d.get("desired_state")))' "$resp" "$TOKEN_FILE"
  chmod 600 "$TOKEN_FILE"
}

handle_desired() {
  local desired resp action
  resp="$(api GET /api/v1/nodes/desired)"
  desired="$(python3 -c 'import json,sys; print(json.loads(sys.argv[1]).get("desired_state",""))' "$resp")"
  echo "Desired state: $desired"
  if [[ "$desired" == "provision" ]]; then
    if [[ ! -f "${STATE_DIR}/provisioned.flag" ]]; then
      echo "Downloading provision script..."
      curl -fsSL "${CONTROL_PLANE_URL}/download/provision-kvm.sh" -o /tmp/nexora-provision-kvm.sh
      chmod +x /tmp/nexora-provision-kvm.sh
      echo "Running KVM/libvirt provision steps (idempotent)..."
      bash /tmp/nexora-provision-kvm.sh || true
      touch "${STATE_DIR}/provisioned.flag"
    else
      echo "Provision already completed (flag present)."
    fi
  fi
}

heartbeat_once() {
  local checks payload
  checks="$(collect_checks)"
  payload="$(python3 -c 'import json,sys; print(json.dumps({"checks": json.loads(sys.argv[1])}))' "$checks")"
  api POST /api/v1/nodes/heartbeat "$payload" >/dev/null
  handle_desired
}

enroll_if_needed

if [[ "${1:-}" == "--once" ]]; then
  heartbeat_once
  exit 0
fi

echo "Starting heartbeat loop every ${INTERVAL}s (Ctrl+C to stop)"
while true; do
  heartbeat_once || echo "Heartbeat failed; will retry"
  sleep "$INTERVAL"
done