#!/usr/bin/env bash
#
# OllamaRedirect agent — one-command Linux installer.
#
#   curl -fsSL https://www.magpy.co/download/install.sh | bash
#
# Installs Ollama + the OllamaRedirect agent (a self-contained static binary),
# pulls two starter models, enrolls the device to your uSaaS workspace, and sets
# it to run either as a `systemd --user` service or manually. Assumes the Nvidia
# driver is already installed.
#
# Config (env; overridable, but the defaults below point at uSaaS production):
#   OR_BROKER_URL   wss endpoint agents dial   (prod: wss://gpu.magpy.co/agent)
#   OR_AUTH_URL     public auth/enroll base     (prod: https://gpu.magpy.co)
#   OR_DL_BASE      where the agent binary is hosted (prod: https://www.magpy.co/download)
#   OR_MODELS       starter models (default: "llama3.1:8b nomic-embed-text")
#   OR_SHA256_<arch> expected agent checksum (baked in below; integrity-checked)
#   OR_DRY_RUN=1    print actions without executing (for review)
set -euo pipefail

OR_BROKER_URL="${OR_BROKER_URL:-wss://gpu.magpy.co/agent}"
OR_AUTH_URL="${OR_AUTH_URL:-https://gpu.magpy.co}"
OR_DL_BASE="${OR_DL_BASE:-https://www.magpy.co/download}"
OR_MODELS="${OR_MODELS:-llama3.1:8b nomic-embed-text}"
OR_DRY_RUN="${OR_DRY_RUN:-0}"

# Integrity: expected SHA-256 of each published agent binary (from scripts/release
# or a plain `CGO_ENABLED=0 go build`). Verified after download; a mismatch aborts.
OR_SHA256_amd64="${OR_SHA256_amd64:-b593951ab95800b0ad01f1a222491f51e610813d90f50b8c58081301c1386941}"
OR_SHA256_arm64="${OR_SHA256_arm64:-cfd44dd78c93a01dd4b66c0a2f46018c6e6b96f0b8ef505fe0aef8e4c40b36a2}"

BIN_DIR="$HOME/.local/bin"
CFG_DIR="$HOME/.config/ollamaredirect"
BIN="$BIN_DIR/ollamaredirect-agent"
ENVFILE="$CFG_DIR/agent.env"
SECRET="$CFG_DIR/secret.json"
UNIT="$HOME/.config/systemd/user/ollamaredirect-agent.service"

c()   { printf '\033[1;36m%s\033[0m\n' "$*"; }        # step
ok()  { printf '\033[1;32m  ✓ %s\033[0m\n' "$*"; }
warn(){ printf '\033[1;33m  ! %s\033[0m\n' "$*"; }
die() { printf '\033[1;31merror: %s\033[0m\n' "$*" >&2; exit 1; }
run() { if [ "$OR_DRY_RUN" = "1" ]; then echo "  [dry-run] $*"; else eval "$@"; fi; }

# Prompts must read from the terminal, not the curl pipe (this script is run as
# `curl … | bash`, so stdin is the download, not the keyboard). Open /dev/tty on
# fd 3; if there's no controlling terminal (headless/CI), fall back to the default.
ask() {
  local p="$1" d="$2" a=""
  if { exec 3<>/dev/tty; } 2>/dev/null; then
    read -r -p "$p" a <&3 || true
    exec 3<&- || true
  fi
  echo "${a:-$d}"
}

# ── 0. Preflight ──────────────────────────────────────────────────────────────
c "OllamaRedirect agent installer"
[ "$(uname -s)" = "Linux" ] || die "this installer is for Linux"
command -v curl >/dev/null || die "curl is required"

case "$(uname -m)" in
  x86_64|amd64) ARCH=amd64 ;;
  aarch64|arm64) ARCH=arm64 ;;
  *) die "unsupported architecture: $(uname -m)" ;;
esac
ok "Linux/$ARCH"

if command -v nvidia-smi >/dev/null 2>&1; then
  ok "Nvidia GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)"
else
  warn "nvidia-smi not found — Ollama will fall back to CPU (slow). Install the Nvidia driver for GPU."
fi

# ── 1. Base deps (python3 per request; not required by the static-binary agent) ─
c "Checking base dependencies"
if command -v python3 >/dev/null 2>&1; then
  ok "python3 present ($(python3 --version 2>&1))"
else
  warn "python3 missing — installing (note: the agent itself does not need it)"
  if   command -v apt-get >/dev/null; then run "sudo apt-get update -qq && sudo apt-get install -y python3"
  elif command -v dnf     >/dev/null; then run "sudo dnf install -y python3"
  elif command -v pacman  >/dev/null; then run "sudo pacman -Sy --noconfirm python"
  else warn "unknown package manager — install python3 manually if you need it"; fi
fi

# ── 2. Ollama ─────────────────────────────────────────────────────────────────
c "Installing Ollama"
if command -v ollama >/dev/null 2>&1; then
  ok "Ollama already installed ($(ollama --version 2>/dev/null | head -1))"
else
  run "curl -fsSL https://ollama.com/install.sh | sh"
  command -v ollama >/dev/null || [ "$OR_DRY_RUN" = "1" ] || die "Ollama install failed"
  ok "Ollama installed"
fi

# ── 2b. GPU backend compatibility (older archs need Ollama's cuda_v12) ─────────
# Ollama auto-selects the cuda_v13 backend, but CUDA 13 dropped Maxwell/Pascal/
# Volta kernels — so GPUs with compute capability < 7.5 (GTX 10-series = 6.1,
# V100 = 7.0) get NO GPU kernels and silently fall back to CPU. Pin cuda_v12 for
# them via a systemd drop-in on Ollama's service. Turing+ (>=7.5) keeps the default.
c "Checking GPU/CUDA backend compatibility"
if command -v nvidia-smi >/dev/null 2>&1; then
  # Lowest compute capability across all GPUs as major*10+minor (6.1 -> 61).
  MINCAP="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \
            | awk -F. 'NF>=2{v=$1*10+$2; if(m==""||v<m)m=v} END{print m}')"
  if [ -z "${MINCAP:-}" ]; then
    warn "couldn't read GPU compute capability (old driver?) — if inference runs on CPU, set OLLAMA_LLM_LIBRARY=cuda_v12 on the ollama service manually"
  elif [ "$MINCAP" -lt 75 ]; then
    CAP_H="$(printf '%s' "$MINCAP" | sed 's/\(.\)$/.\1/')"
    if systemctl show ollama -p Environment 2>/dev/null | grep -q cuda_v12; then
      ok "Ollama already pinned to cuda_v12 (GPU compute $CAP_H < 7.5)"
    else
      DROPIN=/etc/systemd/system/ollama.service.d/10-cuda-v12.conf
      warn "GPU compute $CAP_H < 7.5 — pinning Ollama to cuda_v12 (cuda_v13 has no kernels for it)"
      if [ "$OR_DRY_RUN" = "1" ]; then
        echo "  [dry-run] write $DROPIN (OLLAMA_LLM_LIBRARY=cuda_v12) + systemctl restart ollama"
      else
        sudo mkdir -p "$(dirname "$DROPIN")"
        printf '[Service]\nEnvironment="OLLAMA_LLM_LIBRARY=cuda_v12"\n' | sudo tee "$DROPIN" >/dev/null
        sudo systemctl daemon-reload 2>/dev/null || true
        sudo systemctl restart ollama 2>/dev/null || warn "restart the ollama service so cuda_v12 takes effect"
        # Wait for Ollama to come back so the model-pull step below doesn't race.
        for _ in $(seq 1 30); do curl -fsS http://127.0.0.1:11434/api/tags >/dev/null 2>&1 && break; sleep 1; done
        ok "Ollama pinned to cuda_v12 ($DROPIN)"
      fi
    fi
  else
    ok "GPU compute $((MINCAP/10)).$((MINCAP%10)) — default cuda_v13 backend is fine"
  fi
else
  warn "no nvidia-smi — skipping CUDA backend check (CPU inference)"
fi

# ── 3. Agent binary ───────────────────────────────────────────────────────────
c "Installing the OllamaRedirect agent"
run "mkdir -p '$BIN_DIR' '$CFG_DIR'"
URL="$OR_DL_BASE/agent-linux-$ARCH"
run "curl -fSL '$URL' -o '$BIN.tmp'"
# Optional integrity check against a provided checksum.
eval "EXPECT=\${OR_SHA256_$ARCH:-}"
if [ -n "${EXPECT:-}" ] && [ "$OR_DRY_RUN" != "1" ]; then
  got=$(sha256sum "$BIN.tmp" | awk '{print $1}')
  [ "$got" = "$EXPECT" ] || die "checksum mismatch for agent-linux-$ARCH (got $got)"
  ok "checksum verified"
fi
run "chmod 0755 '$BIN.tmp' && mv '$BIN.tmp' '$BIN'"
[ "$OR_DRY_RUN" = "1" ] || ok "agent: $("$BIN" -version 2>/dev/null | head -1)"

# ── 4. Config ─────────────────────────────────────────────────────────────────
c "Writing config"
AGENT_ID="$(cat /proc/sys/kernel/random/uuid 2>/dev/null || echo "agent-$(hostname)-$$")"
GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo unknown)"
if [ "$OR_DRY_RUN" = "1" ]; then
  echo "  [dry-run] write $ENVFILE (broker=$OR_BROKER_URL auth=$OR_AUTH_URL id=$AGENT_ID)"
else
  umask 077
  cat > "$ENVFILE" <<EOF
OR_BROKER_URL=$OR_BROKER_URL
OR_AUTH_URL=$OR_AUTH_URL
OR_AGENT_ID=$AGENT_ID
OR_SECRET_FILE=$SECRET
OR_GPU=$GPU_NAME
OR_SECRET_BACKEND=file
EOF
  ok "config → $ENVFILE"
fi

# ── 5. Models ─────────────────────────────────────────────────────────────────
c "Pulling starter models: $OR_MODELS"
# Drop a re-runnable puller next to the config so the user can add models later.
if [ "$OR_DRY_RUN" != "1" ]; then
  cat > "$CFG_DIR/init-models.sh" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
MODELS="${OR_MODELS:-llama3.1:8b nomic-embed-text}"
command -v ollama >/dev/null || { echo "ollama not found" >&2; exit 1; }
have() { ollama list 2>/dev/null | awk 'NR>1{print $1}' | grep -qx "$1"; }
for m in $MODELS; do
  if have "$m"; then echo "==> $m already present"; else echo "==> pulling $m…"; ollama pull "$m"; fi
done
EOF
  chmod +x "$CFG_DIR/init-models.sh"
fi
run "OR_MODELS='$OR_MODELS' '$CFG_DIR/init-models.sh'"

# ── 6. Enroll (interactive — prints a code you approve in uSaaS) ───────────────
c "Enrolling this device to your workspace"
if [ "$OR_DRY_RUN" = "1" ]; then
  echo "  [dry-run] $BIN -enroll  (would print a device code to approve in Settings → Local GPU)"
else
  echo "  A code will appear below. Open uSaaS → Settings → Local GPU and enter it."
  set -a; . "$ENVFILE"; set +a
  "$BIN" -enroll || die "enrollment failed"
  ok "device enrolled"
fi

# ── 7. Run mode ───────────────────────────────────────────────────────────────
c "How should the agent run?"
echo "  [S] systemd --user service (auto-start, survives logout)"
echo "  [m] manual (you start it yourself)"
MODE="$(ask '  Choose [S/m]: ' S)"

case "$MODE" in
  m|M|manual)
    ok "Manual mode. Start the agent with:"
    echo "     set -a; . '$ENVFILE'; set +a; '$BIN'"
    ;;
  *)
    c "Installing systemd --user service"
    if [ "$OR_DRY_RUN" = "1" ]; then
      echo "  [dry-run] write $UNIT; systemctl --user enable --now ollamaredirect-agent; loginctl enable-linger $USER"
    else
      mkdir -p "$(dirname "$UNIT")"
      cat > "$UNIT" <<EOF
[Unit]
Description=OllamaRedirect agent (uSaaS GPU tunnel)
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=$BIN
EnvironmentFile=$ENVFILE
Restart=always
RestartSec=5
NoNewPrivileges=true

[Install]
WantedBy=default.target
EOF
      systemctl --user daemon-reload
      systemctl --user enable --now ollamaredirect-agent.service
      # Keep the user service running after logout / across reboots.
      loginctl enable-linger "$USER" 2>/dev/null || warn "could not enable linger (service stops at logout)"
      ok "service running: systemctl --user status ollamaredirect-agent"
    fi
    ;;
esac

c "Done."
echo "  Logs (systemd):  journalctl --user -u ollamaredirect-agent -f"
echo "  Re-pull models:  OR_MODELS='…' $CFG_DIR/init-models.sh"