#!/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
DISPLAY_TIER="$TIER"
PRODUCT_CATEGORY="server"
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"
LICENSE_CHECK_SCRIPT="$SCRIPT_DIR/scripts/cezen-license-check.py"
FEASIBILITY_JSON="/opt/cezen/feasibility.json"
LICENSE_JSON="/opt/cezen/license.json"
OVERRIDE_JSON="/opt/cezen/license.override.json"
INSTALL_RECORD_JSON="/opt/cezen/install-record.json"
PUBLIC_KEY_PATH="${CEZEN_LICENSE_PUBLIC_KEY:-$SCRIPT_DIR/autoinstall/keys/cezen-license-public.pem}"

# 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
    workstation)
      DISPLAY_TIER="workstation"
      PRODUCT_CATEGORY="workstation"
      TIER="starter"
      ;;
    entry|basic)      TIER="basic" ;;
    mid|pro)          TIER="pro" ;;
    advanced|max)     TIER="max" ;;
    starter)          TIER="starter" ;;
  esac
  if [ "$PRODUCT_CATEGORY" = "server" ]; then
    DISPLAY_TIER="$TIER"
  fi
}

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
}

license_eval_field() {
  local expr="$1"
  python3 - "$LICENSE_EVAL_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
}

run_license_evaluation() {
  LICENSE_EVAL_JSON="/tmp/cezen-license-eval.json"
  if [ -f "$LICENSE_CHECK_SCRIPT" ]; then
    python3 "$LICENSE_CHECK_SCRIPT" \
      --license "$LICENSE_JSON" \
      --override "$OVERRIDE_JSON" \
      --feasibility "$FEASIBILITY_JSON" \
      --public-key "$PUBLIC_KEY_PATH" > "$LICENSE_EVAL_JSON"
  else
    echo "ERROR: License checker not found: $LICENSE_CHECK_SCRIPT"
    exit 1
  fi
}

enforce_tier_constraints() {
  run_license_evaluation
  local license_status allowed_tier hardware_tier selected_state
  license_status="$(license_eval_field license.status)"
  allowed_tier="$(license_eval_field license.allowed_tier)"
  hardware_tier="$(license_eval_field hardware.recommended_tier)"
  selected_state="$(python3 - "$LICENSE_EVAL_JSON" "$TIER" <<'PY'
import json, sys
d=json.load(open(sys.argv[1]))
tier=sys.argv[2]
for opt in d.get("tier_options", []):
    if opt.get("tier") == tier:
        print(opt.get("state", "disabled_by_license"))
        break
else:
    print("disabled_by_license")
PY
)"
  echo "→ License status: ${license_status:-missing} | Allowed tier: ${allowed_tier:-basic} | Hardware tier: ${hardware_tier:-starter}"
  case "$selected_state" in
    enabled|override_required)
      ;;
    disabled_by_license)
      echo "ERROR: Selected tier '$TIER' exceeds the current license allowance (${allowed_tier:-basic})."
      exit 1
      ;;
    disabled_by_hardware)
      echo "ERROR: Selected tier '$TIER' exceeds hardware feasibility (${hardware_tier:-starter})."
      exit 1
      ;;
    *)
      echo "ERROR: Selected tier '$TIER' is not allowed (state: $selected_state)."
      exit 1
      ;;
  esac
}

write_install_record() {
  python3 - "$INSTALL_RECORD_JSON" "$TIER" "$PROFILE" "$SKIP_ROLES" "$GPU_AVAILABLE" "$FEASIBILITY_JSON" "$LICENSE_EVAL_JSON" "$DISPLAY_TIER" "$PRODUCT_CATEGORY" <<'PY'
import json, sys
from datetime import datetime, timezone
from pathlib import Path

out = Path(sys.argv[1])
tier = sys.argv[2]
profile = sys.argv[3]
skip_roles = [r for r in sys.argv[4].split(",") if r]
gpu_available = sys.argv[5].lower() == "true"
feasibility = {}
license_eval = {}
for src, dest in ((sys.argv[6], "feasibility"), (sys.argv[7], "license_eval")):
    try:
        with open(src) as fh:
            data = json.load(fh)
        if dest == "feasibility":
            feasibility = data
        else:
            license_eval = data
    except Exception:
        pass
components = {
    "ollama": "ollama" not in skip_roles,
    "jupyterlab": "jupyterlab" not in skip_roles,
    "chromadb": "chromadb" not in skip_roles,
    "vllm": "vllm" not in skip_roles,
    "mlflow": "mlflow" not in skip_roles,
    "minio": "minio" not in skip_roles,
    "monitoring": "monitoring" not in skip_roles,
    "k3s": "k3s" not in skip_roles,
}
payload = {
    "schema": "cezen.install_record.v1",
    "generated_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z"),
    "selected_tier": sys.argv[8],
    "provisioned_tier": tier,
    "category": sys.argv[9],
    "provisioned_profile": profile,
    "skip_roles": skip_roles,
    "components": components,
    "gpu_available": gpu_available,
    "hardware_recommended_tier": ((feasibility.get("recommendation") or {}).get("recommended_tier") or "starter"),
    "hardware_recommended_profile": ((feasibility.get("recommendation") or {}).get("recommended_profile") or "core"),
    "licensed_tier": (((license_eval.get("license") or {}).get("allowed_tier")) or "basic"),
    "license_status": ((license_eval.get("license") or {}).get("status") or "missing"),
    "install_type": ((license_eval.get("license") or {}).get("install_type") or "field-staging"),
    "override_active": bool(((license_eval.get("override") or {}).get("allow_hardware_mismatch"))),
}
out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(json.dumps(payload, indent=2))
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
}

warn_tier_vs_feasibility() {
  [ -f "$FEASIBILITY_JSON" ] || return 0
  local recommended_tier recommended_profile
  recommended_tier="$(json_field recommendation.recommended_tier)"
  recommended_profile="$(json_field recommendation.recommended_profile)"
  [ -n "$recommended_tier" ] || return 0
  if [ "$recommended_tier" != "$TIER" ]; then
    echo "⚠ Feasibility recommends tier '$recommended_tier' / profile '${recommended_profile:-unknown}' for this hardware."
    echo "  Selected tier remains '$TIER'; unsupported services may be skipped by the feasibility profile."
  else
    echo "✓ Feasibility matches selected tier '$TIER' / profile '${recommended_profile:-unknown}'."
  fi
}

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
StandardError=journal

[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 "╚══════════════════════════════════════════╝"

  apply_profile_from_feasibility
  warn_tier_vs_feasibility

  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"
  write_install_record

  # 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

  # Fixed-width box so every row lines up regardless of tier-name length.
  local box_w=62
  pad_row() { printf '║  %-*s║\n' "$((box_w-4))" "$1"; }
  echo ""
  printf '╔%s╗\n' "$(printf '═%.0s' $(seq 1 "$box_w"))"
  pad_row "Nexus One AI installation complete!"
  pad_row "Tier: $DISPLAY_TIER"
  pad_row ""
  pad_row "Portal      -> http://localhost"
  pad_row "Ollama API  -> http://localhost:11434"
  pad_row "vLLM API    -> http://localhost:8000"
  pad_row "JupyterLab  -> http://localhost:8888"
  pad_row "MLflow      -> http://localhost:5000"
  pad_row "Grafana     -> http://localhost:3000"
  pad_row ""
  pad_row "Admin login -> admin / Cezen@2024!"
  pad_row "(you will be required to change this on first login)"
  printf '╚%s╝\n' "$(printf '═%.0s' $(seq 1 "$box_w"))"
}

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

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

check_root
run_feasibility

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

enforce_tier_constraints

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