#!/usr/bin/env bash
set -euo pipefail
ensure_root
detect_pkg_manager
pkg_install curl
install_docker

if ask_to_install "Ollama Router"; then
  echo ""
  read -rp "Listen-Port des Router (Default 11437): " ROUTER_PORT
  ROUTER_PORT=${ROUTER_PORT:-11437}
  echo ""
  read -rp "NVIDIA Node IP: " NVIDIA_IP
  read -rp "NVIDIA Node Port (Default 11436): " NVIDIA_PORT
  NVIDIA_PORT=${NVIDIA_PORT:-11436}
  echo ""
  read -rp "AMD (ROCm) Node IP: " AMD_IP
  read -rp "AMD Node Port (Default 11435): " AMD_PORT
  AMD_PORT=${AMD_PORT:-11435}
  echo ""
  read -rp "CPU-only Node IP: " CPU_IP
  read -rp "CPU Node Port (Default 11434): " CPU_PORT
  CPU_PORT=${CPU_PORT:-11434}
  BASE="/srv/docker/services/ollama-router"
  $SUDO mkdir -p "${BASE}"
  cd "${BASE}"
  $SUDO tee config.yml >/dev/null <<'EOF'
routes:
  llama3.1:8b-instruct:
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  mistral-nemo:12b:
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  huihui_ai/deepseek-r1-abliterated:14b:
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  phi3:medium-128k:
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  mxbai-embed-large:
    - url: http://${CPU_IP}:${CPU_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
  phi3:mini:
    - url: http://${CPU_IP}:${CPU_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
  gemma2:2b-instruct-q6_K:
    - url: http://${CPU_IP}:${CPU_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
  qwen2.5-coder:14b:
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  deepseek-coder-v2:16b:
    - url: http://${AMD_IP}:${AMD_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${CPU_IP}:${CPU_PORT}
  qwen2.5-coder:7b:
    - url: http://${CPU_IP}:${CPU_PORT}
    - url: http://${NVIDIA_IP}:${NVIDIA_PORT}
    - url: http://${AMD_IP}:${AMD_PORT}
EOF
  $SUDO sed -i "s|\${NVIDIA_IP}|${NVIDIA_IP}|g; s|\${NVIDIA_PORT}|${NVIDIA_PORT}|g; s|\${AMD_IP}|${AMD_IP}|g; s|\${AMD_PORT}|${AMD_PORT}|g; s|\${CPU_IP}|${CPU_IP}|g; s|\${CPU_PORT}|${CPU_PORT}|g" config.yml
  $SUDO tee docker-compose.yml >/dev/null <<EOF
version: "3.9"
services:
  ollama-router:
    image: ghcr.io/ollama/ollama-router:latest
    container_name: ollama-router
    restart: unless-stopped
    networks: [ai]
    volumes:
      - ./config.yml:/app/config.yml:ro
    ports:
      - "${ROUTER_PORT}:11437"
networks:
  ai:
    external: true
EOF
  $SUDO docker network inspect ai >/dev/null 2>&1 || $SUDO docker network create ai
  CPU_MODELS=(
    "phi3:mini"
    "gemma2:2b-instruct-q6_K"
    "mxbai-embed-large"
    "qwen2.5-coder:7b"
  )
  for m in "${CPU_MODELS[@]}"; do
    echo "→ Pull ${m} on CPU node ${CPU_IP}:${CPU_PORT}"
    $SUDO curl -fsSL -X POST "http://${CPU_IP}:${CPU_PORT}/api/pull" -d "{"name":"${m}"}" || true
  done
  log "✅ Router konfiguriert in ${BASE}"
  log "ℹ️  Start:  cd ${BASE} && docker compose up -d"
else
  log "${YELLOW}⏭ Ollama Router übersprungen.${NC}"
fi