#!/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)
#   --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
RVS_TOKEN_ARG="${RVS_TOKEN:-}"
RVS_HOST_ARG="${RVS_HOST:-}"
while [[ $# -gt 0 ]]; do
  case "$1" in
    --up) DO_UP=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-Format) ──
step "APT-Komponenten (contrib non-free non-free-firmware)"
add_component() {
  local file="$1" comp="$2"
  # nur in Components:-Zeilen, und nur wenn das Wort noch nicht (als ganzes) da ist
  if ! grep -qE "^Components:.*[[:space:]]${comp}([[:space:]]|\$)" "$file"; then
    sed -i -E "s/^(Components:.*)\$/\1 ${comp}/" "$file"
    return 0
  fi
  return 1
}
CHANGED=0
DEB822="/etc/apt/sources.list.d/debian.sources"
if [[ -f "$DEB822" ]]; then
  for comp in contrib non-free non-free-firmware; do
    if add_component "$DEB822" "$comp"; then CHANGED=1; ok "hinzugefuegt: $comp"; fi
  done
  [[ $CHANGED -eq 0 ]] && ok "bereits aktiv (debian.sources)"
elif [[ -f /etc/apt/sources.list ]]; then
  # Legacy-Format: an die deb-Zeilen anhaengen falls noetig
  if ! grep -qE "non-free-firmware" /etc/apt/sources.list; then
    sed -i -E '/^deb .*debian/ s/$/ contrib non-free non-free-firmware/' /etc/apt/sources.list
    CHANGED=1; ok "hinzugefuegt (sources.list)"
  else
    ok "bereits aktiv (sources.list)"
  fi
else
  warn "Keine APT-Sources gefunden — non-free ggf. manuell aktivieren."
fi
[[ $CHANGED -eq 1 ]] && apt-get update -qq

# ── 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
    warn "Installiere nvidia-driver (zieht Kernel-Modul via DKMS, dauert)…"
    apt-get install -y -qq nvidia-driver firmware-misc-nonfree >/dev/null
    ok "nvidia-driver 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 "Speaker-ID + F5-Voice von der alten Gamebox kopieren (nicht in git):"
warn "   xtts/voice-id/  und  xtts/voices/  →  $XTTS_DIR/"

# ── 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