#!/usr/bin/env bash
#
# ARIA KI-Box (gpubox) Bootstrap — frisches Debian Trixie → startklarer
# GPU-Satelliten-Host fuer den xtts-Stack (Voxtral/Whisper STT, F5-TTS, lokales LLM).
#
# Ablauf nach `git clone`:
#   cd ARIA-AGENT/ai-box
#   sudo ./bootstrap.sh                 # richtet Treiber + Docker + NVIDIA-Toolkit ein
#   # (falls Treiber frisch installiert: einmal `sudo reboot`, dann Script erneut)
#   sudo ./bootstrap.sh --up            # dazu: xtts-Stack (Whisper+F5+LLM) hochziehen
#
# Optionen:
#   --up            am Ende den xtts-Stack starten (Default-Profil, OHNE voxtral)
#   --upgrade-driver  NVIDIA-Treiber aus trixie-backports (modernes CUDA fuer Voxtral).
#                     Danach REBOOT noetig. Default-Bootstrap laesst 550 unangetastet.
#   --token <TOK>   RVS_TOKEN in xtts/.env eintragen (alternativ: env RVS_TOKEN=...)
#   --rvs-host <H>  RVS_HOST setzen (Default aus .env.example)
#
# IDEMPOTENT: bereits erledigte Schritte werden uebersprungen. Nach dem
# Treiber-Reboot einfach nochmal ausfuehren — der Rest laeuft dann durch.
#
set -euo pipefail

# ── CLI ──
DO_UP=0
DO_UPGRADE_DRIVER=0
RVS_TOKEN_ARG="${RVS_TOKEN:-}"
RVS_HOST_ARG="${RVS_HOST:-}"
while [[ $# -gt 0 ]]; do
  case "$1" in
    --up) DO_UP=1; shift ;;
    --upgrade-driver) DO_UPGRADE_DRIVER=1; shift ;;
    --token) RVS_TOKEN_ARG="${2:-}"; shift 2 ;;
    --rvs-host) RVS_HOST_ARG="${2:-}"; shift 2 ;;
    -h|--help) grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) echo "Unbekannte Option: $1"; exit 1 ;;
  esac
done

# ── Log-Helfer ──
c_g="\033[1;32m"; c_y="\033[1;33m"; c_r="\033[1;31m"; c_b="\033[1;34m"; c_0="\033[0m"
STEP=0
step() { STEP=$((STEP+1)); echo -e "\n${c_b}[${STEP}] $*${c_0}"; }
ok()   { echo -e "  ${c_g}✓${c_0} $*"; }
warn() { echo -e "  ${c_y}!${c_0} $*"; }
die()  { echo -e "${c_r}✗ $*${c_0}" >&2; exit 1; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
XTTS_DIR="$REPO_ROOT/xtts"

# ── Preflight ──
[[ "$(id -u)" -eq 0 ]] || die "Bitte als root ausfuehren (sudo ./bootstrap.sh)."
command -v apt-get >/dev/null || die "Kein apt-get — dieses Script ist fuer Debian/Trixie."
export DEBIAN_FRONTEND=noninteractive

echo -e "${c_b}=== ARIA KI-Box Bootstrap ===${c_0}"
echo "Repo:  $REPO_ROOT"
echo "xtts:  $XTTS_DIR"

# ── 1. Basis-Pakete ──
step "Basis-Pakete"
apt-get update -qq
apt-get install -y -qq ca-certificates curl gnupg git lsb-release >/dev/null
ok "ca-certificates, curl, gnupg, git"

# ── 2. non-free Repos aktivieren (Trixie deb822 + Legacy) ──
step "APT-Komponenten (contrib non-free non-free-firmware)"
# Ergaenzt die drei Komponenten in JEDER Components:-Zeile einer deb822-Datei —
# pro Zeile nur was fehlt, reihenfolge-robust, idempotent. Die Adressen pruefen
# ganze Woerter (non-free-firmware zaehlt NICHT als non-free).
add_components() {
  sed -i -E '/^[Cc]omponents:/{
    /(^|[[:space:]])contrib([[:space:]]|$)/!s/$/ contrib/
    /(^|[[:space:]])non-free([[:space:]]|$)/!s/$/ non-free/
    /(^|[[:space:]])non-free-firmware([[:space:]]|$)/!s/$/ non-free-firmware/
  }' "$1"
}
ENABLED_ANY=0
shopt -s nullglob
for f in /etc/apt/sources.list.d/*.sources; do
  grep -qE '^[Cc]omponents:' "$f" || continue
  b="$(md5sum "$f")"; add_components "$f"; a="$(md5sum "$f")"
  if [[ "$b" != "$a" ]]; then ENABLED_ANY=1; ok "aktualisiert: $(basename "$f")"; fi
done
shopt -u nullglob
# Legacy /etc/apt/sources.list (deb-Zeilen)
if [[ -f /etc/apt/sources.list ]] && grep -qE '^deb ' /etc/apt/sources.list; then
  if ! grep -qE '^deb .*[[:space:]]non-free([[:space:]]|$)' /etc/apt/sources.list; then
    sed -i -E '/^deb .*debian/ s/$/ contrib non-free non-free-firmware/' /etc/apt/sources.list
    ENABLED_ANY=1; ok "aktualisiert: sources.list"
  fi
fi
[[ $ENABLED_ANY -eq 0 ]] && ok "non-free schon aktiv"
apt-get update -qq   # immer neu einlesen, damit der Kandidat sicher da ist

# ── 2b. (opt-in) Treiber-Upgrade via trixie-backports — fuer Voxtral/modernes CUDA ──
# Der Trixie-Standardtreiber (550, CUDA 12.4) ist zu alt fuer den modernen Stack
# (torch 2.13, vLLM). Backports bringt einen neueren, apt-verwalteten Treiber.
# NUR mit --upgrade-driver, damit ein laufendes Setup nicht ungewollt angefasst wird.
if [[ "$DO_UPGRADE_DRIVER" -eq 1 ]]; then
  step "NVIDIA-Treiber-Upgrade (trixie-backports)"
  BP_FILE="/etc/apt/sources.list.d/backports.sources"
  if ! grep -rqs "trixie-backports" /etc/apt/sources.list /etc/apt/sources.list.d/ 2>/dev/null; then
    cat > "$BP_FILE" <<'EOF'
Types: deb
URIs: http://deb.debian.org/debian
Suites: trixie-backports
Components: main contrib non-free non-free-firmware
Signed-By: /usr/share/keyrings/debian-archive-keyring.gpg
EOF
    ok "trixie-backports hinzugefuegt"
  else
    ok "trixie-backports bereits aktiv"
  fi
  apt-get update
  apt-get install -y linux-headers-amd64 || true
  apt-get install -y "linux-headers-$(uname -r)" || true
  if apt-get install -y -t trixie-backports nvidia-driver; then
    dkms autoinstall >/dev/null 2>&1 || true
    NEWV="$(dpkg-query -W -f='${Version}' nvidia-driver 2>/dev/null || echo '?')"
    ok "nvidia-driver aus backports installiert: ${NEWV}"
    echo
    echo -e "${c_y}==> REBOOT noetig, damit der neue Treiber laedt:${c_0}"
    echo -e "${c_y}    sudo reboot   &&   danach:  cd ai-box && sudo ./bootstrap.sh --up --token <TOKEN>${c_0}"
    echo -e "${c_y}    (nvidia-smi zeigt dann die neue Version + CUDA-Level)${c_0}"
    exit 0
  else
    warn "backports-Install fehlgeschlagen — 550er bleibt aktiv."
    warn "Alternative fuer den neuesten Treiber: NVIDIAs CUDA-Repo fuer debian13"
    warn "  (developer.download.nvidia.com/compute/cuda/repos/debian13/x86_64) → Paket 'cuda-drivers'."
    die "Treiber-Upgrade nicht moeglich — siehe oben."
  fi
fi

# ── 3. NVIDIA-Treiber ──
step "NVIDIA-Treiber"
if nvidia-smi >/dev/null 2>&1; then
  ok "Treiber aktiv: $(nvidia-smi --query-gpu=name --format=csv,noheader | paste -sd', ')"
  DRIVER_ACTIVE=1
else
  if dpkg -l | grep -q '^ii  nvidia-driver '; then
    warn "Treiber installiert, aber nvidia-smi antwortet nicht → REBOOT noetig."
    DRIVER_ACTIVE=0
  else
    # KEIN separater Kandidaten-Check — die apt-cache-Ausgabe ist locale-/pipe-
    # fragil (hat faelschlich "kein Kandidat" gemeldet). Der Install IST der Test:
    # direkt installieren; schlaegt er fehl, einmal volles apt-get update + Retry,
    # dann erst mit Diagnose abbrechen.
    # Kernel-Header ZUERST — sonst ueberspringt DKMS den Modulbau ("No kernel
    # headers were found") und nvidia-smi kann spaeter nicht mit dem Treiber reden.
    warn "Kernel-Header + nvidia-driver installieren (DKMS-Build, dauert)…"
    apt-get install -y linux-headers-amd64 || true
    apt-get install -y "linux-headers-$(uname -r)" || true
    if ! apt-get install -y nvidia-driver firmware-misc-nonfree; then
      warn "Install fehlgeschlagen — volles 'apt-get update' + zweiter Versuch…"
      apt-get update
      if ! apt-get install -y nvidia-driver firmware-misc-nonfree; then
        echo
        warn "nvidia-driver liess sich nicht installieren. Aktive Quellen:"
        grep -rhE '^[Cc]omponents:' /etc/apt/sources.list.d/*.sources 2>/dev/null | sed 's/^/     /' || true
        grep -E '^deb ' /etc/apt/sources.list 2>/dev/null | sed 's/^/     /' || true
        die "Pruefe 'apt-cache policy nvidia-driver' + Netz/non-free."
      fi
    fi
    # Falls nvidia-kernel-dkms schon (ohne Header) installiert war: Modul jetzt bauen.
    dkms autoinstall >/dev/null 2>&1 || true
    ok "nvidia-driver + Kernel-Modul installiert"
    DRIVER_ACTIVE=0
  fi
fi

# ── 4. Docker Engine + Compose-Plugin ──
step "Docker"
if command -v docker >/dev/null 2>&1; then
  ok "Docker vorhanden: $(docker --version)"
else
  warn "Installiere Docker (offizielles get.docker.com)…"
  curl -fsSL https://get.docker.com | sh >/dev/null
  systemctl enable --now docker >/dev/null 2>&1 || true
  ok "Docker installiert: $(docker --version)"
fi
if docker compose version >/dev/null 2>&1; then
  ok "Compose-Plugin: $(docker compose version | head -1)"
else
  warn "Compose-Plugin fehlt — installiere docker-compose-plugin…"
  apt-get install -y -qq docker-compose-plugin >/dev/null || \
    warn "Konnte docker-compose-plugin nicht via apt holen — get.docker.com bringt es normalerweise mit."
fi

# ── 5. NVIDIA Container Toolkit ──
step "NVIDIA Container Toolkit"
NCT_LIST="/etc/apt/sources.list.d/nvidia-container-toolkit.list"
NCT_KEY="/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
if ! dpkg -l | grep -q '^ii  nvidia-container-toolkit '; then
  [[ -f "$NCT_KEY" ]] || curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
    | gpg --dearmor -o "$NCT_KEY"
  if [[ ! -f "$NCT_LIST" ]]; then
    curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
      | sed "s#deb https://#deb [signed-by=${NCT_KEY}] https://#g" > "$NCT_LIST"
  fi
  apt-get update -qq
  apt-get install -y -qq nvidia-container-toolkit >/dev/null
  ok "nvidia-container-toolkit installiert"
else
  ok "nvidia-container-toolkit vorhanden"
fi
# Docker-Runtime auf NVIDIA konfigurieren (idempotent)
if ! grep -q '"nvidia"' /etc/docker/daemon.json 2>/dev/null; then
  nvidia-ctk runtime configure --runtime=docker >/dev/null
  systemctl restart docker
  ok "Docker-Runtime auf NVIDIA konfiguriert + Docker neugestartet"
else
  ok "Docker-Runtime bereits NVIDIA-konfiguriert"
fi

# ── 6. xtts/.env vorbereiten ──
step "xtts/.env"
if [[ -f "$XTTS_DIR/.env" ]]; then
  ok ".env existiert bereits (unangetastet)"
else
  cp "$XTTS_DIR/.env.example" "$XTTS_DIR/.env"
  ok ".env aus .env.example erstellt"
fi
if [[ -n "$RVS_TOKEN_ARG" ]]; then
  sed -i -E "s#^RVS_TOKEN=.*#RVS_TOKEN=${RVS_TOKEN_ARG}#" "$XTTS_DIR/.env"
  ok "RVS_TOKEN eingetragen"
fi
if [[ -n "$RVS_HOST_ARG" ]]; then
  sed -i -E "s#^RVS_HOST=.*#RVS_HOST=${RVS_HOST_ARG}#" "$XTTS_DIR/.env"
  ok "RVS_HOST=${RVS_HOST_ARG} eingetragen"
fi
if grep -q '^RVS_TOKEN=dein_token_hier' "$XTTS_DIR/.env"; then
  warn "RVS_TOKEN ist noch der Platzhalter — vor dem Start setzen:"
  warn "   nano $XTTS_DIR/.env   (oder: sudo ./bootstrap.sh --token <TOKEN>)"
fi
warn "Stimm-Daten (nicht in git): falls von der alten Box noch vorhanden, xtts/voice-id/"
warn "   + xtts/voices/ herkopieren. Sonst egal — in der App neu anlegen:"
warn "   Sprache neu enrollen (Speaker-ID, sonst fail-open) + F5-Referenz neu hochladen."

# ── 7. GPU-im-Container verifizieren ──
step "GPU-im-Container Test"
if [[ "${DRIVER_ACTIVE:-0}" -eq 1 ]]; then
  if docker run --rm --gpus all nvidia/cuda:12.4.0-base-ubuntu22.04 nvidia-smi >/dev/null 2>&1; then
    ok "Docker sieht die GPU — KI-Box ist einsatzbereit."
    GPU_READY=1
  else
    warn "Host-Treiber ok, aber Container sieht die GPU nicht — Toolkit/Runtime pruefen."
    GPU_READY=0
  fi
else
  warn "Treiber noch nicht aktiv → Test uebersprungen."
  echo
  echo -e "${c_y}==> REBOOT noetig, dann Script erneut ausfuehren:${c_0}"
  echo -e "${c_y}    sudo reboot   &&   (nach dem Boot)   sudo ./bootstrap.sh${c_0}"
  exit 0
fi

# ── 8. Optional: xtts-Stack hochziehen ──
if [[ $DO_UP -eq 1 && "${GPU_READY:-0}" -eq 1 ]]; then
  step "xtts-Stack starten (Default-Profil: Whisper + F5 + LLM — OHNE voxtral)"
  warn "Erster Start laedt Modelle (mehrere GB via HuggingFace) — kann dauern."
  ( cd "$XTTS_DIR" && docker compose up -d --build )
  ok "Stack laeuft. Logs:  docker logs -f aria-whisper-bridge"
  echo
  echo "  voxtral (STT via Voxtral) braucht >=16 GB VRAM → erst mit der 24-GB-Karte:"
  echo "     cd $XTTS_DIR && docker compose stop whisper-bridge && docker compose --profile voxtral up -d --build"
fi

# ── Abschluss ──
echo
echo -e "${c_g}=== Fertig. KI-Box startklar. ===${c_0}"
if [[ $DO_UP -eq 0 ]]; then
  echo "Naechster Schritt — Stack starten:"
  echo "   cd $XTTS_DIR && docker compose up -d --build"
  echo "oder direkt:  sudo ./bootstrap.sh --up"
fi