#!/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"
FEASIBILITY_SCRIPT="$AIPACKAGE_DIR/scripts/cezen-feasibility.sh"
LICENSE_CHECK_SCRIPT="$AIPACKAGE_DIR/scripts/cezen-license-check.py"
PUBLIC_KEY_PATH="${CEZEN_LICENSE_PUBLIC_KEY:-$AIPACKAGE_DIR/autoinstall/keys/cezen-license-public.pem}"
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}')}"

# ── Nexus One AI Workstation detection ─────────────────────
# Workstation is a standalone product category, not a rung on the Server
# S/M/L/Max ladder (see cezen_license.WORKSTATION_TIER). It must never be
# offered as a choice in the Server tier-selection menu below (Step 3).
# When this ISO was built as a Workstation image (CEZEN_TIER=workstation,
# or the /opt/cezen/tier marker says "workstation"), Step 3 is skipped
# entirely and the existing "starter" Ansible profile is reused for
# provisioning — same compact single-GPU stack, different commercial
# category. TIER="starter" here only selects which Ansible playbook runs;
# it does not change the "workstation" marker written to /opt/cezen/tier
# or the install-record tier fields below.
WORKSTATION_MODE=false
if [ "${CEZEN_TIER:-}" = "workstation" ]; then
  WORKSTATION_MODE=true
elif [ -f /opt/cezen/tier ] && [ "$(tr -d '[:space:]' < /opt/cezen/tier)" = "workstation" ]; then
  WORKSTATION_MODE=true
fi

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_PATH=$(whiptail --title "$TITLE" \
  --inputbox "\nStep 2 of 4: License & Customer Details\n\nPath to signed license JSON:\n\nLeave blank for offline field staging.\nExample: /media/usb/customer-license.json" \
  $H $W "" 3>&1 1>&2 2>&3)

OVERRIDE_PATH=$(whiptail --title "$TITLE" \
  --inputbox "\nOptional: Cezen hardware override JSON path\n\nLeave blank unless support explicitly provided one." \
  $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\nOptional local record field for field installs." \
  $H $W "" 3>&1 1>&2 2>&3)

mkdir -p /opt/cezen
rm -f /opt/cezen/license.json /opt/cezen/license.override.json

if [ -n "$LICENSE_PATH" ]; then
  if [ ! -f "$LICENSE_PATH" ]; then
    whiptail --title "$TITLE" \
      --msgbox "\nLicense file not found:\n  $LICENSE_PATH" \
      $H $W
    exit 1
  fi
  cp "$LICENSE_PATH" /opt/cezen/license.json
  chmod 0640 /opt/cezen/license.json
  chown root:cezen /opt/cezen/license.json 2>/dev/null || true
fi

if [ -n "$OVERRIDE_PATH" ]; then
  if [ ! -f "$OVERRIDE_PATH" ]; then
    whiptail --title "$TITLE" \
      --msgbox "\nOverride file not found:\n  $OVERRIDE_PATH" \
      $H $W
    exit 1
  fi
  cp "$OVERRIDE_PATH" /opt/cezen/license.override.json
  chmod 0640 /opt/cezen/license.override.json
  chown root:cezen /opt/cezen/license.override.json 2>/dev/null || true
fi

bash "$FEASIBILITY_SCRIPT" /opt/cezen/feasibility.json >> "$LOG_FILE" 2>&1
python3 "$LICENSE_CHECK_SCRIPT" \
  --license /opt/cezen/license.json \
  --override /opt/cezen/license.override.json \
  --feasibility /opt/cezen/feasibility.json \
  --public-key "$PUBLIC_KEY_PATH" > /tmp/cezen-license-check.json

LICENSE_STATUS=$(python3 - <<'PY'
import json
d=json.load(open("/tmp/cezen-license-check.json"))
print((d.get("license") or {}).get("status","missing"))
PY
)
# Human-readable label for display only — the raw code above (e.g.
# "invalid_signature", "machine_mismatch") is what gets persisted to
# install-record.json and is meant for machines/support, not end users.
LICENSE_STATUS_LABEL=$(LICENSE_STATUS="$LICENSE_STATUS" python3 - <<'PY'
import os
labels = {
    "valid":             "Valid",
    "missing":           "No license installed (field-staging mode)",
    "expired":           "Expired — contact support@cezentech.com",
    "invalid_signature": "Invalid signature — re-upload a valid signed license",
    "not_yet_valid":     "Not yet in its valid date range",
    "machine_mismatch":  "Bound to different hardware — contact support@cezentech.com",
}
status = os.environ.get("LICENSE_STATUS", "missing")
print(labels.get(status, status))
PY
)
LICENSE_ALLOWED_TIER=$(python3 - <<'PY'
import json
d=json.load(open("/tmp/cezen-license-check.json"))
print((d.get("license") or {}).get("allowed_tier","basic"))
PY
)
HARDWARE_TIER=$(python3 - <<'PY'
import json
d=json.load(open("/tmp/cezen-license-check.json"))
print((d.get("hardware") or {}).get("recommended_tier","starter"))
PY
)

# ════════════════════════════════════════════════════════════
# STEP 3: SELECT TIER
# (skipped entirely for Nexus One AI Workstation — it is not a Server
#  tier and must never appear in, or be selected from, this menu)
# ════════════════════════════════════════════════════════════

if [ "$WORKSTATION_MODE" = true ]; then
  TIER="starter"   # Ansible/install.sh routing only — reuses the compact
                    # single-GPU profile. The "workstation" marker in
                    # /opt/cezen/tier and the install-record below is what
                    # actually drives Workstation branding/entitlement.
  whiptail --title "$TITLE" \
    --msgbox "\nNexus One AI Workstation\n\nLicense status: ${LICENSE_STATUS_LABEL}\nHardware recommendation: ${HARDWARE_TIER}\n\nThis image installs the personal-appliance stack (local chat, personal RAG, document intelligence, prompt studio) — it is not a Server tier and will not be offered a Server S/M/L/Max upgrade path." \
    $H $W
else
  mapfile -t TIER_MENU < <(python3 - <<'PY'
import json
d=json.load(open("/tmp/cezen-license-check.json"))
labels = {
    "starter": "Server S — 1x RTX 5090 / 32GB VRAM · Small team",
    "basic":   "Server M — 1x NVIDIA RTX Pro 6000 (96GB) · Up to 20 users",
    "pro":     "Server L — 2x RTX 5090 / RTX Pro class · Up to 100 users",
    "max":     "Server Max — 4-8x H100/H200/A100 class · 100+ users",
}
for opt in d.get("tier_options", []):
    if opt.get("selectable"):
        print(opt["tier"])
        print(labels.get(opt["tier"], opt["tier"]))
PY
  )

  if [ "${#TIER_MENU[@]}" -eq 0 ]; then
    whiptail --title "$TITLE" \
      --msgbox "\nNo installable tiers are available.\n\nLicense status: ${LICENSE_STATUS_LABEL}\nHardware recommendation: ${HARDWARE_TIER}\n\nCheck the signed license or contact Cezen support." \
      $H $W
    exit 1
  fi

  whiptail --title "$TITLE" \
    --msgbox "\nLicense status: ${LICENSE_STATUS_LABEL}\nAllowed tier: ${LICENSE_ALLOWED_TIER}\nHardware recommendation: ${HARDWARE_TIER}\n\nOnly valid tiers will be shown next." \
    $H $W

  TIER=$(whiptail --title "$TITLE" \
    --menu "\nStep 3 of 4: Select AI Package Tier\n\nChoose the tier allowed by your license and hardware:" \
    $H $W 4 \
    "${TIER_MENU[@]}" \
    3>&1 1>&2 2>&3)
fi

# Display-only label — $TIER itself stays the Ansible/install.sh routing
# value (e.g. "starter"); TIER_DISPLAY is what the operator sees and what
# gets recorded as the product category for Workstation installs.
if [ "$WORKSTATION_MODE" = true ]; then
  TIER_DISPLAY="workstation"
else
  TIER_DISPLAY="$TIER"
fi

# ════════════════════════════════════════════════════════════
# 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_PATH" ]; then
  LICENSE_DISPLAY="Signed file (${LICENSE_STATUS_LABEL})"
fi

whiptail --title "$TITLE" \
  --yesno "\nReady to install. Please confirm:\n\nNetwork:  ${NET_MODE} (${MY_IP})\nCustomer: ${CUSTOMER_NAME:-Not entered}\nLicense:  ${LICENSE_DISPLAY}\nAllowed:  ${LICENSE_ALLOWED_TIER}\nHardware: ${HARDWARE_TIER}\nTier:     ${TIER_DISPLAY}\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 SUPPORT_UNTIL TIER TIER_DISPLAY LICENSE_STATUS LICENSE_ALLOWED_TIER HARDWARE_TIER WORKSTATION_MODE
python3 - <<'PY'
import json, os, time

# NOTE: install.sh (Phase 2) rewrites /opt/cezen/install-record.json after
# this. For Workstation, it preserves category="workstation" and selected_tier
# = "workstation", while provisioned_tier becomes "starter" to reflect the
# Ansible profile actually installed. The persistent Workstation marker is
# still /opt/cezen/tier, and install.sh never touches that marker.
payload = {
    "schema": "cezen.install_record.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(),
    "selected_tier": os.environ.get("TIER_DISPLAY", "basic").strip(),
    "provisioned_tier": os.environ.get("TIER_DISPLAY", "basic").strip(),
    "category": "workstation" if os.environ.get("WORKSTATION_MODE") == "true" else "server",
    "support_until": os.environ.get("SUPPORT_UNTIL", "").strip(),
    "install_type": "licensed" if os.path.exists("/opt/cezen/license.json") else "field-staging",
    "license_status": os.environ.get("LICENSE_STATUS", "missing").strip(),
    "licensed_tier": os.environ.get("LICENSE_ALLOWED_TIER", "basic").strip(),
    "hardware_recommended_tier": os.environ.get("HARDWARE_TIER", "starter").strip(),
    "captured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}

with open("/opt/cezen/install-record.json", "w") as f:
    json.dump(payload, f, indent=2)
PY
chown root:cezen /opt/cezen/install-record.json 2>/dev/null || true
chmod 0640 /opt/cezen/install-record.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      -> http://localhost\nOllama API  -> http://localhost:11434\nGrafana     -> http://localhost:3000\n\nAdmin login -> admin / Cezen@2024!\n(you will be required to change this on first login)\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