#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────
# Nexus One AI — First Boot Setup Wizard
# Runs on first boot after OS install via systemd service.
# Uses whiptail for the TUI.
# ─────────────────────────────────────────────────────────────
set -e

AIPACKAGE_DIR="/opt/aipackage"
LOG_FILE="/var/log/cezen-setup.log"
INSTALL_LOG_FILE="/var/log/cezen-install.log"
export TERM="${TERM:-linux}"
touch "$LOG_FILE" "$INSTALL_LOG_FILE"
chmod 0644 "$LOG_FILE" "$INSTALL_LOG_FILE"

detect_iface() {
  ip route show default 2>/dev/null | awk '/default/ {print $5; exit}'
}

IFACE="$(detect_iface)"
IFACE="${IFACE:-$(ip -o link show | awk -F': ' '$2 !~ /lo|docker|br-|veth/ {print $2; exit}')}"

netmask_to_prefix() {
  case "$1" in
    32|31|30|29|28|27|26|25|24|23|22|21|20|19|18|17|16|15|14|13|12|11|10|9|8|7|6|5|4|3|2|1) echo "$1" ;;
    255.255.255.255) echo 32 ;;
    255.255.255.254) echo 31 ;;
    255.255.255.252) echo 30 ;;
    255.255.255.248) echo 29 ;;
    255.255.255.240) echo 28 ;;
    255.255.255.224) echo 27 ;;
    255.255.255.192) echo 26 ;;
    255.255.255.128) echo 25 ;;
    255.255.255.0)   echo 24 ;;
    255.255.254.0)   echo 23 ;;
    255.255.252.0)   echo 22 ;;
    255.255.248.0)   echo 21 ;;
    255.255.240.0)   echo 20 ;;
    255.255.224.0)   echo 19 ;;
    255.255.192.0)   echo 18 ;;
    255.255.128.0)   echo 17 ;;
    255.255.0.0)     echo 16 ;;
    *) return 1 ;;
  esac
}

# ── Colors / terminal setup ────────────────────────────────
export NEWT_COLORS='
root=,black
window=black,white
border=white,black
title=black,white
button=black,white
actbutton=white,blue
checkbox=black,white
actcheckbox=white,blue
entry=black,white
label=black,white
listbox=black,white
actlistbox=white,blue
textbox=black,white
acttextbox=white,blue
'

TITLE="  Nexus One AI — Server Setup  "
H=20
W=70

# ── Welcome ────────────────────────────────────────────────
whiptail --title "$TITLE" \
  --msgbox "\nWelcome to the Nexus One AI installer.\n\nThis wizard will configure your network, capture license/customer details, and install the AI stack.\n\nMake sure this server is connected to the internet before continuing." \
  $H $W

# ════════════════════════════════════════════════════════════
# STEP 1: NETWORK CONFIGURATION
# ════════════════════════════════════════════════════════════

NET_MODE=$(whiptail --title "$TITLE" \
  --menu "\nStep 1 of 4: Network Configuration\n\nHow should this server get its IP address?" \
  $H $W 2 \
  "DHCP"   "Automatic (get IP from network)" \
  "Static" "Manual (enter IP address)" \
  3>&1 1>&2 2>&3)

if [ "$NET_MODE" = "Static" ]; then
  mkdir -p /etc/cloud/cloud.cfg.d
  printf "network: {config: disabled}\n" > /etc/cloud/cloud.cfg.d/99-disable-network-config.cfg
  rm -f /etc/netplan/50-cloud-init.yaml /etc/netplan/99-cezen-dhcp.yaml || true

  IP_ADDR=$(whiptail --title "$TITLE" \
    --inputbox "\nEnter static IP address:\n(Example: 192.168.1.100)" \
    $H $W "" 3>&1 1>&2 2>&3)

  NETMASK=$(whiptail --title "$TITLE" \
    --inputbox "\nEnter subnet mask / CIDR prefix:\n(Examples: 24 or 255.255.255.0)" \
    $H $W "24" 3>&1 1>&2 2>&3)

  PREFIX="$(netmask_to_prefix "$NETMASK")" || {
    whiptail --title "$TITLE" \
      --msgbox "\nInvalid subnet mask or CIDR prefix:\n  ${NETMASK}\n\nUse a CIDR prefix such as 24, or a mask such as 255.255.255.0." \
      $H $W
    exit 1
  }

  GATEWAY=$(whiptail --title "$TITLE" \
    --inputbox "\nEnter default gateway:\n(Example: 192.168.1.1)" \
    $H $W "" 3>&1 1>&2 2>&3)

  DNS=$(whiptail --title "$TITLE" \
    --inputbox "\nEnter DNS server:\n(Example: 8.8.8.8)" \
    $H $W "8.8.8.8" 3>&1 1>&2 2>&3)

  # Write netplan static config
  cat > /etc/netplan/99-cezen-static.yaml << EOF
network:
  version: 2
  ethernets:
    ${IFACE}:
      dhcp4: false
      addresses:
        - ${IP_ADDR}/${PREFIX}
      routes:
        - to: default
          via: ${GATEWAY}
      nameservers:
        addresses: [${DNS}]
EOF
  chmod 0600 /etc/netplan/99-cezen-static.yaml

  if ! netplan apply; then
    whiptail --title "$TITLE" \
      --msgbox "\nFailed to apply network settings.\n\nIP:      ${IP_ADDR}/${PREFIX}\nGateway: ${GATEWAY}\nDNS:     ${DNS}\n\nPlease verify the address, gateway, and subnet." \
      $H $W
    exit 1
  fi
  sleep 3

  # Verify connectivity
  if ! ping -c 2 -W 3 8.8.8.8 &>/dev/null; then
    whiptail --title "$TITLE" \
      --msgbox "\nWARNING: Cannot reach internet with these settings.\n\nIP:      ${IP_ADDR}/${PREFIX}\nGateway: ${GATEWAY}\nDNS:     ${DNS}\n\nPlease verify your network settings. The install requires internet access." \
      $H $W
  else
    whiptail --title "$TITLE" \
      --msgbox "\nNetwork configured successfully!\n\nIP:      ${IP_ADDR}/${PREFIX}\nGateway: ${GATEWAY}\nDNS:     ${DNS}" \
      $H $W
  fi

else
  # DHCP — persist it explicitly and confirm it's working
  rm -f /etc/netplan/99-cezen-static.yaml || true
  cat > /etc/netplan/99-cezen-dhcp.yaml << EOF
network:
  version: 2
  ethernets:
    ${IFACE}:
      dhcp4: true
EOF

  netplan apply 2>/dev/null || true
  sleep 2
  MY_IP=$(hostname -I | awk '{print $1}')
  if [ -n "$MY_IP" ]; then
    whiptail --title "$TITLE" \
      --msgbox "\nDHCP configured successfully!\n\nServer IP: ${MY_IP}\n\nYou can switch to a static IP later by editing:\n  /etc/netplan/99-cezen-dhcp.yaml" \
      $H $W
  else
    whiptail --title "$TITLE" \
      --msgbox "\nWARNING: No IP address obtained via DHCP.\n\nMake sure the network cable is connected and try rebooting." \
      $H $W
  fi
fi

# ════════════════════════════════════════════════════════════
# STEP 2: LICENSE / CUSTOMER DETAILS
# ════════════════════════════════════════════════════════════

CUSTOMER_NAME=$(whiptail --title "$TITLE" \
  --inputbox "\nStep 2 of 4: License & Customer Details\n\nCustomer / organisation name:" \
  $H $W "" 3>&1 1>&2 2>&3)

CUSTOMER_ID=$(whiptail --title "$TITLE" \
  --inputbox "\nStep 2 of 4: License & Customer Details\n\nCustomer ID / project code:\n\nLeave blank for evaluation or field staging." \
  $H $W "" 3>&1 1>&2 2>&3)

CONTACT_EMAIL=$(whiptail --title "$TITLE" \
  --inputbox "\nStep 2 of 4: License & Customer Details\n\nCustomer/admin contact email:\n\nLeave blank if unavailable on site." \
  $H $W "" 3>&1 1>&2 2>&3)

LICENSE_KEY=$(whiptail --title "$TITLE" \
  --passwordbox "\nStep 2 of 4: License & Customer Details\n\nLicense key / activation code:\n\nLeave blank for offline field staging. Models and keys can be added later." \
  $H $W "" 3>&1 1>&2 2>&3)

SUPPORT_UNTIL=$(whiptail --title "$TITLE" \
  --inputbox "\nStep 2 of 4: License & Customer Details\n\nSupport valid until (YYYY-MM-DD):\n\nLeave blank if not issued yet." \
  $H $W "" 3>&1 1>&2 2>&3)

# ════════════════════════════════════════════════════════════
# STEP 3: SELECT TIER
# ════════════════════════════════════════════════════════════

TIER=$(whiptail --title "$TITLE" \
  --menu "\nStep 3 of 4: Select AI Package Tier\n\nChoose the tier that matches your hardware:" \
  $H $W 4 \
  "starter" "Starter — 1× RTX 5090 / 32GB VRAM · Small team" \
  "basic" "Entry — 1× NVIDIA RTX Pro 6000 (96GB) · Up to 20 users" \
  "pro"   "Pro   — 2× RTX 5090 / RTX Pro class · Up to 100 users" \
  "max"   "Max   — 4–8× H100/H200/A100 class · 100+ users" \
  3>&1 1>&2 2>&3)

# ════════════════════════════════════════════════════════════
# STEP 4: SELECT AI TOOLS
# ════════════════════════════════════════════════════════════

TOOLS=$(whiptail --title "$TITLE" \
  --checklist "\nStep 4 of 4: Select AI Tools to Install\n\nSpace to toggle, Enter to confirm:" \
  $H $W 8 \
  "ollama"     "Ollama + Open WebUI  (LLM inference + chat UI)" ON \
  "jupyterlab" "JupyterLab           (Notebook environment)"    ON \
  "chromadb"   "ChromaDB             (Vector DB for RAG)"       ON \
  "vllm"       "vLLM                 (OpenAI-compatible API)"   ON \
  "mlflow"     "MLflow               (Experiment tracking)"     ON \
  "minio"      "MinIO                (S3 model storage)"        ON \
  "monitoring" "Grafana + Prometheus (GPU & system monitoring)" ON \
  "k3s"        "K3s                  (Lightweight Kubernetes)"  ON \
  3>&1 1>&2 2>&3)

# ════════════════════════════════════════════════════════════
# CONFIRM
# ════════════════════════════════════════════════════════════

# Format tools list for display
TOOLS_DISPLAY=$(echo "$TOOLS" | tr -d '"' | tr ' ' '\n' | sed 's/^/  · /' | tr '\n' '\n')
MY_IP=$(hostname -I | awk '{print $1}')
LICENSE_DISPLAY="Field staging / evaluation"
if [ -n "$LICENSE_KEY" ]; then
  LICENSE_DISPLAY="Provided"
fi

whiptail --title "$TITLE" \
  --yesno "\nReady to install. Please confirm:\n\nNetwork:  ${NET_MODE} (${MY_IP})\nCustomer: ${CUSTOMER_NAME:-Not entered}\nLicense:  ${LICENSE_DISPLAY}\nTier:     ${TIER}\n\nTools:\n${TOOLS_DISPLAY}\n\nThis will take 20–40 minutes.\nThe server will reboot once during install (NVIDIA drivers).\n\nContinue?" \
  $H $W

# ════════════════════════════════════════════════════════════
# RUN INSTALLER
# ════════════════════════════════════════════════════════════

clear
echo ""
echo "╔══════════════════════════════════════════╗"
echo "║   Nexus One AI — Installing...         ║"
echo "║   Check progress: journalctl -f           ║"
echo "╚══════════════════════════════════════════╝"
echo ""

# Write selected tools to a config file so install.sh can read it
mkdir -p /opt/cezen
SKIP_ROLES=""
for role in ollama jupyterlab chromadb vllm mlflow minio monitoring k3s; do
  if ! echo "$TOOLS" | grep -q "$role"; then
    if [ -n "$SKIP_ROLES" ]; then
      SKIP_ROLES="${SKIP_ROLES},${role}"
    else
      SKIP_ROLES="${role}"
    fi
  fi
done

cat > /opt/cezen/install.conf << EOF
TIER=${TIER}
SKIP_ROLES=${SKIP_ROLES}
EOF

export CUSTOMER_NAME CUSTOMER_ID CONTACT_EMAIL LICENSE_KEY SUPPORT_UNTIL TIER
python3 - <<'PY'
import json, os, time

payload = {
    "schema": "cezen.license.v1",
    "customer_name": os.environ.get("CUSTOMER_NAME", "").strip(),
    "customer_id": os.environ.get("CUSTOMER_ID", "").strip(),
    "contact_email": os.environ.get("CONTACT_EMAIL", "").strip(),
    "license_key": os.environ.get("LICENSE_KEY", "").strip(),
    "tier": os.environ.get("TIER", "basic").strip(),
    "support_until": os.environ.get("SUPPORT_UNTIL", "").strip(),
    "install_type": "licensed" if os.environ.get("LICENSE_KEY", "").strip() else "field-staging",
    "issued_by": "Cezen",
    "captured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    "signature_status": "unsigned-field-install",
}

with open("/opt/cezen/license.json", "w") as f:
    json.dump(payload, f, indent=2)
PY
chown root:cezen /opt/cezen/license.json 2>/dev/null || true
chmod 0640 /opt/cezen/license.json

whiptail --title "$TITLE" \
  --infobox "\nInstalling Nexus One AI stack...\n\nThis can take several minutes.\n\nLogs are being written to:\n  $INSTALL_LOG_FILE" \
  $H $W

if bash "$AIPACKAGE_DIR/install.sh" --tier="$TIER" >> "$INSTALL_LOG_FILE" 2>&1; then
  # Mark as configured only after the installer finishes successfully.
  touch /opt/cezen/.setup-done
  whiptail --title "$TITLE" \
    --msgbox "\nInstaller command finished successfully.\n\nPortal:\n  http://localhost\n\nFor detailed logs, run:\n  sudo tail -f $INSTALL_LOG_FILE" \
    $H $W
else
  whiptail --title "$TITLE" \
    --msgbox "\nInstaller command failed.\n\nThe setup wizard will run again on next boot.\n\nCheck the log with:\n  sudo tail -n 120 $INSTALL_LOG_FILE" \
    $H $W
  exit 1
fi