#!/usr/bin/env bash
# Reproducible bootstrap for a fresh Debian 12/13 AI host.
set -Eeuo pipefail
umask 077

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG=""
STACK_DIR=/opt/mike-ai/stack
SECRETS_DIR=/etc/mike-ai
STATE_DIR=/srv/mike-ai

log() { printf '\n==> %s\n' "$*"; }
die() { printf 'FEHLER: %s\n' "$*" >&2; exit 1; }

usage() {
  cat <<'EOF'
Aufruf: sudo ./install.sh --config /root/mike-ai-install.env

Das Skript ist idempotent und darf nach einem Treiber-Reboot erneut ausgeführt
werden. Es formatiert keine Datenträger und löscht keine Modelle oder Volumes.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --config) CONFIG="${2:-}"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) die "Unbekanntes Argument: $1" ;;
  esac
done

[[ $EUID -eq 0 ]] || die "Bitte als root ausführen."
[[ -n "$CONFIG" && -f "$CONFIG" ]] || die "Konfigurationsdatei fehlt (--config)."
config_mode=$(stat -c '%a' "$CONFIG")
if (( (8#$config_mode & 077) != 0 )); then
  die "Konfigurationsdatei darf nicht für Gruppe/Andere schreib- oder lesbar sein (chmod 600)."
fi
# shellcheck disable=SC1090
source "$CONFIG"

required=(AI_HOSTNAME ADMIN_USER AI_BIND_ADDRESS MODEL_DIR FAST_MODEL_FILE
  FAST_MODEL_URL FAST_MODEL_SHA256 MEDIUM_MODEL_FILE MEDIUM_MODEL_URL
  MEDIUM_MODEL_SHA256 LONG_MODEL_FILE LONG_MODEL_URL LONG_MODEL_SHA256
  EXPERIMENTAL_MODEL_FILE EXPERIMENTAL_MODEL_URL EXPERIMENTAL_MODEL_SHA256
  VISION_PROJECTOR_FILE VISION_PROJECTOR_URL VISION_PROJECTOR_SHA256)
for name in "${required[@]}"; do
  [[ -n "${!name:-}" ]] || die "Pflichtwert $name fehlt."
done

source /etc/os-release
[[ ${ID:-} == debian ]] || die "Unterstützt wird Debian, gefunden: ${ID:-unbekannt}."
[[ ${VERSION_ID%%.*} == 12 || ${VERSION_ID%%.*} == 13 ]] || \
  die "Unterstützt werden Debian 12 und 13."
[[ $(dpkg --print-architecture) == amd64 ]] || die "Dieser Stack erwartet amd64."
id "$ADMIN_USER" >/dev/null 2>&1 || die "ADMIN_USER $ADMIN_USER existiert nicht."

install_base_packages() {
  log "Basispakete installieren"
  apt-get update
  DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    ca-certificates curl git gnupg jq openssl wireguard-tools iptables \
    iproute2 pciutils rsync unattended-upgrades
}

install_docker() {
  log "Docker CE aus dem offiziellen Repository installieren"
  install -m 0755 -d /etc/apt/keyrings
  curl -fsSL https://download.docker.com/linux/debian/gpg \
    -o /etc/apt/keyrings/docker.asc
  chmod a+r /etc/apt/keyrings/docker.asc
  cat >/etc/apt/sources.list.d/docker.sources <<EOF
Types: deb
URIs: https://download.docker.com/linux/debian
Suites: ${VERSION_CODENAME}
Components: stable
Architectures: $(dpkg --print-architecture)
Signed-By: /etc/apt/keyrings/docker.asc
EOF
  apt-get update
  DEBIAN_FRONTEND=noninteractive apt-get install -y \
    docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
  systemctl enable --now docker
}

install_nvidia() {
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    [[ ${INSTALL_NVIDIA_DRIVER:-false} == true ]] || \
      die "NVIDIA-Treiber fehlt. Installiere ihn oder setze INSTALL_NVIDIA_DRIVER=true."
    log "Aktuellen NVIDIA-Treiber aus dem offiziellen NVIDIA-Repository installieren"
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
      "linux-headers-$(uname -r)"

    local distro="debian${VERSION_ID%%.*}"
    local keyring_url="https://developer.download.nvidia.com/compute/cuda/repos/${distro}/x86_64/cuda-keyring_1.1-1_all.deb"
    local keyring_deb
    keyring_deb=$(mktemp /tmp/cuda-keyring.XXXXXX.deb)
    curl --fail --location --retry 5 --output "$keyring_deb" "$keyring_url"
    dpkg -i "$keyring_deb"
    rm -f "$keyring_deb"
    apt-get update

    if [[ -n ${NVIDIA_DRIVER_BRANCH:-} ]]; then
      local pinning="nvidia-driver-pinning-${NVIDIA_DRIVER_BRANCH}"
      apt-cache show "$pinning" >/dev/null 2>&1 || \
        die "Gewuenschter NVIDIA-Treiberzweig fehlt im offiziellen Repository: $pinning"
      DEBIAN_FRONTEND=noninteractive apt-get install -y "$pinning"
    fi

    # RTX 30/50 werden beide von NVIDIAs offenen Kernelmodulen unterstuetzt.
    # Compute-only vermeidet X11-, Desktop- und Vulkan-Pakete auf dem Headless-Host.
    DEBIAN_FRONTEND=noninteractive apt-get install -y \
      nvidia-driver-cuda nvidia-kernel-open-dkms
    printf '\nNVIDIA-Treiber installiert. Jetzt neu starten und danach denselben Installer erneut ausführen.\n'
    exit 20
  fi
  nvidia-smi >/dev/null || die "nvidia-smi ist vorhanden, aber der Treiber funktioniert nicht."

  local driver_major min_driver_major
  driver_major=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 | cut -d. -f1)
  min_driver_major=${NVIDIA_MIN_DRIVER_MAJOR:-570}
  [[ $driver_major =~ ^[0-9]+$ ]] || die "NVIDIA-Treiberversion konnte nicht gelesen werden."
  (( driver_major >= min_driver_major )) || \
    die "NVIDIA-Treiber $driver_major ist zu alt; mindestens $min_driver_major wird verlangt."

  log "NVIDIA Container Toolkit installieren"
  curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
  curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
    >/etc/apt/sources.list.d/nvidia-container-toolkit.list
  apt-get update
  DEBIAN_FRONTEND=noninteractive apt-get install -y nvidia-container-toolkit
  nvidia-ctk runtime configure --runtime=docker
  systemctl restart docker
  docker run --rm --gpus all nvidia/cuda:12.8.1-base-ubuntu24.04 nvidia-smi >/dev/null
}

setup_wireguard() {
  [[ ${WIREGUARD_ENABLE:-false} == true ]] || return 0
  for name in WG_INTERFACE WG_ADDRESS WG_HOME_SUBNET WG_PEER_PUBLIC_KEY WG_PEER_ENDPOINT; do
    [[ -n "${!name:-}" && ${!name} != REPLACE_* ]] || die "WireGuard-Wert $name fehlt."
  done
  log "WireGuard als ausgehenden Heimnetz-Tunnel konfigurieren"
  install -d -m 0700 /etc/wireguard
  local key_file="/etc/wireguard/${WG_INTERFACE}.key"
  [[ -s $key_file ]] || wg genkey >"$key_file"
  chmod 0600 "$key_file"
  local private_key
  private_key=$(<"$key_file")
  local allowed_ips="$WG_HOME_SUBNET"
  [[ ${WG_ROUTE_AI_INTERNET:-true} == true ]] && allowed_ips="0.0.0.0/0, ::/0"
  {
    printf '[Interface]\nAddress = %s\nPrivateKey = %s\nTable = off\n' "$WG_ADDRESS" "$private_key"
    printf '\n[Peer]\nPublicKey = %s\nEndpoint = %s\nAllowedIPs = %s\nPersistentKeepalive = 25\n' \
      "$WG_PEER_PUBLIC_KEY" "$WG_PEER_ENDPOINT" "$allowed_ips"
    if [[ -n ${WG_PEER_PRESHARED_KEY_FILE:-} ]]; then
      [[ -s $WG_PEER_PRESHARED_KEY_FILE ]] || die "WireGuard-Preshared-Key-Datei fehlt."
      printf 'PresharedKey = %s\n' "$(<"$WG_PEER_PRESHARED_KEY_FILE")"
    fi
  } >"/etc/wireguard/${WG_INTERFACE}.conf"
  chmod 0600 "/etc/wireguard/${WG_INTERFACE}.conf"
  systemctl enable --now "wg-quick@${WG_INTERFACE}"
  ip address show "$WG_INTERFACE" >/dev/null
  local wg_host=${WG_ADDRESS%/*}
  [[ $AI_BIND_ADDRESS == "$wg_host" ]] || \
    die "AI_BIND_ADDRESS muss der WireGuard-Adresse ohne Präfix entsprechen ($wg_host)."
  printf 'WireGuard-Public-Key des KI-Hosts: %s\n' "$(wg pubkey <"$key_file")"
}

install_stack_files() {
  log "Stackdateien installieren"
  install -d -m 0755 "$STACK_DIR" "$MODEL_DIR" "$STATE_DIR/backups"
  rsync -a --delete --exclude .git --exclude '*.local.*' \
    --exclude config/install.env "$ROOT_DIR/" "$STACK_DIR/"
  install -d -m 0700 "$SECRETS_DIR"
  [[ -s $SECRETS_DIR/router-api-key ]] || openssl rand -base64 48 >$SECRETS_DIR/router-api-key
  [[ -s $SECRETS_DIR/controller-token ]] || openssl rand -base64 48 >$SECRETS_DIR/controller-token
  [[ -s $SECRETS_DIR/webui-secret ]] || openssl rand -base64 48 >$SECRETS_DIR/webui-secret
  chmod 0600 "$SECRETS_DIR"/*

  local searx="$STACK_DIR/platform/web-search/searxng-settings.yml"
  if [[ ! -s $searx ]]; then
    cp "$STACK_DIR/platform/web-search/searxng-settings.example.yml" "$searx"
    sed -i "s/CHANGE_ME_GENERATE_RANDOM_SECRET/$(openssl rand -hex 32)/" "$searx"
  fi
  # The official image reads this as its unprivileged uid (977).
  chown root:977 "$searx"
  chmod 0640 "$searx"

  cat >$SECRETS_DIR/stack.env <<EOF
AI_BIND_ADDRESS=$AI_BIND_ADDRESS
MODEL_DIR=$MODEL_DIR
ROUTER_API_KEY=$(<$SECRETS_DIR/router-api-key)
CONTROLLER_TOKEN=$(<$SECRETS_DIR/controller-token)
WEBUI_SECRET_KEY=$(<$SECRETS_DIR/webui-secret)
OPENWEBUI_IMAGE=${OPENWEBUI_IMAGE:-ghcr.io/open-webui/open-webui:v0.9.5}
OPENWEBUI_ENABLE_SIGNUP=${OPENWEBUI_ENABLE_SIGNUP:-false}
AI_DNS=${WG_DNS:-1.1.1.1}
FAST_MODEL_FILE=$FAST_MODEL_FILE
MEDIUM_MODEL_FILE=$MEDIUM_MODEL_FILE
LONG_MODEL_FILE=$LONG_MODEL_FILE
EXPERIMENTAL_MODEL_FILE=$EXPERIMENTAL_MODEL_FILE
VISION_PROJECTOR_FILE=$VISION_PROJECTOR_FILE
FAST_CONTEXT=${FAST_CONTEXT:-76800}
MEDIUM_CONTEXT=${MEDIUM_CONTEXT:-94208}
LONG_CONTEXT=${LONG_CONTEXT:-131072}
EXPERIMENTAL_CONTEXT=${EXPERIMENTAL_CONTEXT:-76800}
FAST_GPU_DEVICES=${TEXT_GPU_DEVICES:-0}
MEDIUM_GPU_DEVICES=${TEXT_GPU_DEVICES:-0}
LONG_GPU_DEVICES=${TEXT_GPU_DEVICES:-0}
EXPERIMENTAL_GPU_DEVICES=${TEXT_GPU_DEVICES:-0}
LLAMA_THREADS=${LLAMA_THREADS:-6}
LLAMA_THREADS_BATCH=${LLAMA_THREADS_BATCH:-6}
EOF
  chmod 0600 $SECRETS_DIR/stack.env
}

download_one() {
  local relative=$1 url=$2 expected=$3 target="$MODEL_DIR/$1"
  install -d -m 0755 "$(dirname "$target")"
  if [[ -f $target ]] && printf '%s  %s\n' "$expected" "$target" | sha256sum -c - >/dev/null 2>&1; then
    printf 'Vorhanden und geprüft: %s\n' "$relative"
    return
  fi
  log "Modell laden: $relative"
  curl --fail --location --continue-at - --retry 5 --retry-all-errors \
    --output "$target.partial" "$url"
  printf '%s  %s\n' "$expected" "$target.partial" | sha256sum -c -
  mv "$target.partial" "$target"
}

download_models() {
  [[ ${SKIP_MODEL_DOWNLOADS:-false} == true ]] && return 0
  declare -A seen=()
  local row relative url hash
  while IFS='|' read -r relative url hash; do
    [[ -n ${seen[$relative]:-} ]] && continue
    seen[$relative]=1
    download_one "$relative" "$url" "$hash"
  done <<EOF
$FAST_MODEL_FILE|$FAST_MODEL_URL|$FAST_MODEL_SHA256
$MEDIUM_MODEL_FILE|$MEDIUM_MODEL_URL|$MEDIUM_MODEL_SHA256
$LONG_MODEL_FILE|$LONG_MODEL_URL|$LONG_MODEL_SHA256
$EXPERIMENTAL_MODEL_FILE|$EXPERIMENTAL_MODEL_URL|$EXPERIMENTAL_MODEL_SHA256
$VISION_PROJECTOR_FILE|$VISION_PROJECTOR_URL|$VISION_PROJECTOR_SHA256
EOF
}

install_routing_guard() {
  [[ ${WIREGUARD_ENABLE:-false} == true ]] || return 0
  log "KI-Container auf WireGuard routen und Uni-Netz als Transit sperren"
  cat >/usr/local/sbin/mike-ai-network-guard <<EOF
#!/usr/bin/env bash
set -euo pipefail
WG=$WG_INTERFACE
HOME_NET=$WG_HOME_SUBNET
TABLE=51820
for NET in 172.30.10.0/24 172.30.30.0/24 172.30.40.0/24 172.30.50.0/24; do
  ip rule add from \"\$NET\" table \"\$TABLE\" priority 12000 2>/dev/null || true
done
ip route replace \"\$HOME_NET\" dev \"\$WG\"
ip route replace \"\$HOME_NET\" dev \"\$WG\" table \"\$TABLE\"
ip route replace blackhole default metric 32767 table \"\$TABLE\"
EOF
  if [[ ${WG_ROUTE_AI_INTERNET:-true} == true ]]; then
    printf 'ip route replace default dev "$WG" metric 10 table "$TABLE"\n' >>/usr/local/sbin/mike-ai-network-guard
  fi
  cat >>/usr/local/sbin/mike-ai-network-guard <<'EOF'
# Explicit forwarding policy: AI containers may use wg0; WireGuard clients may
# reach the two published UI/API ports.  Neither side may use this host as a
# general university<->home router.  The blackhole route above prevents a
# fail-open to the university gateway when wg0 loses its peer/default route.
add_rule() { iptables -C DOCKER-USER "$@" 2>/dev/null || iptables -I DOCKER-USER 1 "$@"; }
add_rule -o "$WG" -j DROP
add_rule -i "$WG" -j DROP
add_rule -s 172.30.0.0/16 -o "$WG" -j ACCEPT
add_rule -i "$WG" -d 172.30.10.0/24 -p tcp -m multiport --dports 8080,8081 -j ACCEPT
add_rule -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
EOF
  chmod 0755 /usr/local/sbin/mike-ai-network-guard
  cat >/etc/systemd/system/mike-ai-network-guard.service <<EOF
[Unit]
Description=Routing guard for Mike AI Docker networks
After=docker.service wg-quick@${WG_INTERFACE}.service
Requires=docker.service wg-quick@${WG_INTERFACE}.service

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/mike-ai-network-guard
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
  systemctl daemon-reload
  systemctl enable --now mike-ai-network-guard.service
}

build_and_start() {
  log "llama.cpp und Plattform-Container bauen"
  local commit
  commit=$(<"$STACK_DIR/platform/llama/LLAMA_CPP_COMMIT")
  cd "$STACK_DIR"
  docker build --build-arg LLAMA_CPP_COMMIT="$commit" \
    -f platform/docker/llama-cpp/Dockerfile -t mike-ai/llama.cpp:local .
  # Creates the shared internal tools network before Open WebUI is created.
  # Web search always starts; HA/ARR/Unraid only start when their root-only
  # secret files and required local artifacts are present.
  "$STACK_DIR/platform/mcp/install-tools.sh"
  docker compose --env-file "$SECRETS_DIR/stack.env" --profile inference create \
    llama-fast llama-medium llama-long llama-experimental
  docker compose --env-file "$SECRETS_DIR/stack.env" up -d --build \
    profile-controller router searxng open-webui

  log "Fast-Profil aktivieren und Readiness prüfen"
  docker compose --env-file "$SECRETS_DIR/stack.env" exec -T router python - <<'PY'
import json, os, time, urllib.request
key = os.environ["ROUTER_API_KEY"]
request = urllib.request.Request(
    "http://127.0.0.1:8081/fast", method="POST",
    headers={"Authorization": f"Bearer {key}"})
with urllib.request.urlopen(request, timeout=700) as response:
    print(json.dumps(json.load(response), indent=2))
deadline = time.monotonic() + 60
while time.monotonic() < deadline:
    try:
        with urllib.request.urlopen("http://127.0.0.1:8081/ready", timeout=5) as response:
            if response.status == 200:
                print("Router und Fast-Profil sind bereit.")
                break
    except Exception:
        pass
    time.sleep(2)
else:
    raise SystemExit("Readiness-Check fehlgeschlagen")
PY
}

hostnamectl set-hostname "$AI_HOSTNAME"
install_base_packages
install_docker
install_nvidia
setup_wireguard
install_stack_files
download_models
install_routing_guard
build_and_start

log "Installation abgeschlossen"
cat <<EOF
OpenWebUI:  http://${AI_BIND_ADDRESS}:8080
Router-API: http://${AI_BIND_ADDRESS}:8081

Die geheimen Schlüssel liegen ausschließlich unter $SECRETS_DIR (0600).
Konfiguriere auf der Heimseite für den KI-Peer mindestens die Rückroute
${WG_ADDRESS:-10.77.0.2/32}; bei Internet-Routing zusätzlich NAT/Forwarding.
EOF