#!/usr/bin/env bash
# ─────────────────────────────────────────────
# Nexus One AI — Installer
# Usage:
#   sudo bash install.sh                       → auto-detect tier, Phase 1
#   sudo bash install.sh --tier=starter        → Starter tier, Phase 1
#   sudo bash install.sh --tier=basic          → Basic tier, Phase 1
#   sudo bash install.sh --tier=pro            → Pro tier, Phase 1
#   sudo bash install.sh --tier=max            → Max tier, Phase 1
#   sudo bash install.sh --phase=2 --tier=...  → Phase 2 only (post-reboot)
#   sudo bash install.sh --software-only       → install on customer-owned hardware
#   sudo bash install.sh --feasibility-only    → scan hardware and exit
#   sudo bash install.sh --skip-model-pull     → install Ollama without preloading models
# ─────────────────────────────────────────────
set -e

# Auto-detect tier from ISO marker written by autoinstall user-data
if [ -f /opt/cezen/tier ]; then
  TIER="$(cat /opt/cezen/tier | tr -d '[:space:]')"
elif [ -f /opt/aipackage/autoinstall/.tier ]; then
  TIER="$(cat /opt/aipackage/autoinstall/.tier | tr -d '[:space:]')"
else
  TIER="basic"    # default if no marker found
fi
PHASE="1"
SKIP_ROLES=""
SOFTWARE_ONLY=false
FEASIBILITY_ONLY=false
SKIP_MODEL_PULL=false
PROFILE="auto"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ANSIBLE_DIR="$SCRIPT_DIR/ansible"
FEASIBILITY_SCRIPT="$SCRIPT_DIR/scripts/cezen-feasibility.sh"
FEASIBILITY_JSON="/opt/cezen/feasibility.json"

# Load saved config (written by web setup UI before phase 1)
[ -f /opt/cezen/install.conf ] && source /opt/cezen/install.conf

for arg in "$@"; do
  case $arg in
    --tier=*)  TIER="${arg#*=}" ;;
    --phase=*) PHASE="${arg#*=}" ;;
    --skip=*)  SKIP_ROLES="${arg#*=}" ;;
    --profile=*) PROFILE="${arg#*=}" ;;
    --software-only) SOFTWARE_ONLY=true ;;
    --feasibility-only) FEASIBILITY_ONLY=true ;;
    --skip-model-pull) SKIP_MODEL_PULL=true ;;
  esac
done

normalize_tier() {
  case "$TIER" in
    entry|basic)      TIER="basic" ;;
    mid|pro)          TIER="pro" ;;
    advanced|max)     TIER="max" ;;
    starter)          TIER="starter" ;;
  esac
}

normalize_tier

# ── Preflight ──────────────────────────────────
check_root() {
  if [ "$EUID" -ne 0 ]; then
    echo "ERROR: Run as root: sudo bash install.sh"
    exit 1
  fi
}

check_os() {
  if [ -f /etc/os-release ]; then
    . /etc/os-release
    if [[ "$ID" != "ubuntu" ]]; then
      echo "ERROR: Ubuntu 22.04 required. Detected: $PRETTY_NAME"
      exit 1
    fi
    echo "✓ OS: $PRETTY_NAME"
  fi
}

install_ansible() {
  if ! command -v ansible-playbook &>/dev/null; then
    echo "→ Installing Ansible..."
    apt-get update -qq
    apt-get install -y -qq ansible python3-pip
  fi
  echo "✓ Ansible ready"
}

append_skip_role() {
  local role="$1"
  if [ -z "$SKIP_ROLES" ]; then
    SKIP_ROLES="$role"
  elif [[ ",$SKIP_ROLES," != *",$role,"* ]]; then
    SKIP_ROLES="$SKIP_ROLES,$role"
  fi
}

run_feasibility() {
  if [ -f "$FEASIBILITY_SCRIPT" ]; then
    bash "$FEASIBILITY_SCRIPT" "$FEASIBILITY_JSON"
  else
    echo "WARNING: Feasibility checker not found: $FEASIBILITY_SCRIPT"
  fi
}

json_field() {
  local expr="$1"
  python3 - "$FEASIBILITY_JSON" "$expr" <<'PY'
import json, sys
try:
    d=json.load(open(sys.argv[1]))
    cur=d
    for part in sys.argv[2].split("."):
        cur=cur[part]
    print(cur)
except Exception:
    print("")
PY
}

apply_profile_from_feasibility() {
  [ -f "$FEASIBILITY_JSON" ] || return 0
  local detected_profile
  detected_profile="$(json_field recommendation.recommended_profile)"
  if [ "$PROFILE" = "auto" ] && [ -n "$detected_profile" ]; then
    PROFILE="$detected_profile"
  fi

  case "$PROFILE" in
    core)
      append_skip_role docker
      append_skip_role k3s
      append_skip_role ollama
      append_skip_role vllm
      append_skip_role jupyterlab
      append_skip_role chromadb
      append_skip_role mlflow
      append_skip_role minio
      append_skip_role monitoring
      SKIP_MODEL_PULL=true
      ;;
    cpu-ai)
      append_skip_role k3s
      append_skip_role vllm
      append_skip_role mlflow
      append_skip_role minio
      SKIP_MODEL_PULL=true
      ;;
    gpu-lite|gpu-starter)
      append_skip_role k3s
      append_skip_role mlflow
      append_skip_role minio
      SKIP_MODEL_PULL=true
      ;;
    gpu-standard)
      append_skip_role mlflow
      append_skip_role minio
      ;;
    gpu-pro|gpu-max)
      ;;
    *)
      echo "WARNING: Unknown profile '$PROFILE'; using explicit skip list only."
      ;;
  esac
}

has_nvidia_pci_gpu() {
  for vendor_file in /sys/bus/pci/devices/*/vendor; do
    [ -f "$vendor_file" ] || continue
    if [ "$(tr '[:upper:]' '[:lower:]' < "$vendor_file")" = "0x10de" ]; then
      return 0
    fi
  done
  return 1
}

has_working_nvidia_driver() {
  command -v nvidia-smi &>/dev/null && nvidia-smi &>/dev/null
}

# ── Phase 1: NVIDIA drivers only ──────────────
run_phase1() {
  echo ""
  echo "╔══════════════════════════════════════════╗"
  echo "║   Nexus One AI — Phase 1: NVIDIA       ║"
  echo "╚══════════════════════════════════════════╝"

  if ! has_nvidia_pci_gpu; then
    echo "No NVIDIA GPU found. Continuing with CPU/non-GPU installation path."
    PHASE="2"
    run_phase2
    return
  fi

  ANSIBLE_STDOUT_CALLBACK=yaml \
  ansible-playbook -i localhost, -c local "$ANSIBLE_DIR/phase1_nvidia.yml" \
    -e "tier=$TIER"

  # Register phase 2 as a one-shot systemd service so it runs after reboot
  cat > /etc/systemd/system/cezen-phase2.service << EOF
[Unit]
Description=Nexus One AI Phase 2 Installer
After=network-online.target nvidia-persistenced.service
Wants=network-online.target

[Service]
Type=oneshot
ExecStart=/bin/bash -lc 'set -o pipefail; /bin/bash ${SCRIPT_DIR}/install.sh --phase=2 --tier=${TIER} 2>&1 | tee -a /var/log/cezen-install.log'
RemainAfterExit=yes
StandardOutput=journal+console
StandardError=journal+console

[Install]
WantedBy=multi-user.target
EOF

  systemctl daemon-reload
  systemctl enable cezen-phase2.service

  echo ""
  echo "✓ Phase 2 registered — will run automatically after reboot"
  echo "→ Rebooting in 10 seconds..."
  sleep 10
  reboot
}

# ── Phase 2: Full stack ────────────────────────
run_phase2() {
  echo ""
  echo "╔══════════════════════════════════════════╗"
  echo "║   Nexus One AI — Phase 2: Stack        ║"
  echo "╚══════════════════════════════════════════╝"

  GPU_AVAILABLE=false
  if ! has_working_nvidia_driver; then
    echo "No working NVIDIA GPU/driver found. Continuing with CPU/non-GPU installation path."
    echo "GPU-only features such as NVIDIA Docker runtime, DCGM metrics, and vLLM serving will be skipped or left inactive."
  else
    GPU_AVAILABLE=true
    echo "✓ NVIDIA driver: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"
  fi

  # Build skip_roles extra var (comma-separated list, empty string = skip nothing)
  EXTRA_VARS="tier=$TIER skip_roles=\"$SKIP_ROLES\" gpu_available=$GPU_AVAILABLE skip_model_pull=$SKIP_MODEL_PULL"
  echo "→ Tier: $TIER | Skip: ${SKIP_ROLES:-none}"
  echo "→ GPU available: $GPU_AVAILABLE"
  echo "→ Skip model pull: $SKIP_MODEL_PULL"

  # Select Ansible playbook by tier
  case "$TIER" in
    starter)        PLAYBOOK="$ANSIBLE_DIR/starter.yml" ;;
    basic|entry)    PLAYBOOK="$ANSIBLE_DIR/entry.yml" ;;
    pro)            PLAYBOOK="$ANSIBLE_DIR/pro.yml" ;;
    max)            PLAYBOOK="$ANSIBLE_DIR/max.yml" ;;
    *)
      echo "ERROR: Unknown tier '$TIER'. Valid: starter | basic | pro | max"
      exit 1
      ;;
  esac

  echo "→ Playbook: $PLAYBOOK"
  ANSIBLE_STDOUT_CALLBACK=yaml \
  ansible-playbook -i localhost, -c local "$PLAYBOOK" \
    -e "$EXTRA_VARS"

  # Disable one-shot service so it doesn't run again on next reboot
  systemctl disable cezen-phase2.service 2>/dev/null || true

  echo ""
  echo "╔══════════════════════════════════════════╗"
  echo "║   Nexus One AI installation complete!  ║"
  echo "║   Tier: $(printf '%-33s' "$TIER")║"
  echo "║                                          ║"
  echo "║   Portal     → http://localhost          ║"
  echo "║   Ollama API → http://localhost:11434    ║"
  echo "║   vLLM API   → http://localhost:8000     ║"
  echo "║   Grafana    → http://localhost:3000     ║"
  echo "╚══════════════════════════════════════════╝"
}

# ── Main ───────────────────────────────────────
check_os

if [ "$FEASIBILITY_ONLY" = true ]; then
  run_feasibility
  exit 0
fi

check_root
run_feasibility

if [ "$SOFTWARE_ONLY" = true ]; then
  PHASE="2"
  apply_profile_from_feasibility
fi

install_ansible

if [ "$PHASE" = "1" ]; then
  run_phase1
elif [ "$PHASE" = "2" ]; then
  run_phase2
else
  echo "ERROR: Unknown phase '$PHASE'. Use --phase=1 or --phase=2"
  exit 1
fi