#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────
# Cezen AI Suite — 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"
exec > >(tee -a "$LOG_FILE") 2>&1

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

TITLE="  Cezen AI Suite — Server Setup  "
H=20
W=70

# ── Welcome ────────────────────────────────────────────────
whiptail --title "$TITLE" \
  --msgbox "\nWelcome to the Cezen AI Suite installer.\n\nThis wizard will configure your network 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 3: 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

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

  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(Example: 24  means 255.255.255.0)" \
    $H $W "24" 3>&1 1>&2 2>&3)

  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}/${NETMASK}
      routes:
        - to: default
          via: ${GATEWAY}
      nameservers:
        addresses: [${DNS}]
EOF

  netplan apply 2>/dev/null || true
  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}/${NETMASK}\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}/${NETMASK}\nGateway: ${GATEWAY}\nDNS:     ${DNS}" \
      $H $W
  fi

else
  # DHCP — just confirm it's working
  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 set a static IP later via:\n  sudo nano /etc/netplan/50-cloud-init.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: SELECT TIER
# ════════════════════════════════════════════════════════════

TIER=$(whiptail --title "$TITLE" \
  --menu "\nStep 2 of 3: Select AI Package Tier\n\nChoose the tier that matches your hardware:" \
  $H $W 3 \
  "entry"    "Entry    — 3× NVIDIA L40S   (48GB each) · Up to 20 users" \
  "mid"      "Mid      — RTX Pro 6000 BW  (96GB each) · Up to 50 users" \
  "advanced" "Advanced — HGX H200         (141GB each) · 200+ users" \
  3>&1 1>&2 2>&3)

# ════════════════════════════════════════════════════════════
# STEP 3: SELECT AI TOOLS
# ════════════════════════════════════════════════════════════

TOOLS=$(whiptail --title "$TITLE" \
  --checklist "\nStep 3 of 3: 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}')

whiptail --title "$TITLE" \
  --yesno "\nReady to install. Please confirm:\n\nNetwork:  ${NET_MODE} (${MY_IP})\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 "║   Cezen AI Suite — 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

# Mark as configured so this wizard doesn't run again
touch /opt/cezen/.setup-done

# Run the installer
bash "$AIPACKAGE_DIR/install.sh" --tier="$TIER"