#!/usr/bin/env bash
#
# EthoVision worker — one-step installer for macOS and Linux.
#
# What this does, in order:
#   1. Checks for conda; installs Miniconda into ~/miniconda3 if missing.
#   2. Creates (or reuses) a conda env named "ethovision" with Python 3.11.
#   3. Detects your GPU and installs the matching PyTorch wheels.
#   4. Installs DeepLabCut 3.0 with the SuperAnimal model zoo.
#   5. Downloads the worker script itself into the current directory.
#   6. Prints the exact command to start the worker.
#
# Run with:
#   bash install.sh
#
# Optional environment variables:
#   ETHOVISION_SERVER  — base URL of your EthoVision server
#   ETHOVISION_TOKEN   — worker bearer token
# If set, these are written into a run.sh launcher so you don't have to pass them each time.

set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_NAME="${ETHOVISION_ENV:-ethovision}"
PY_VER="3.11"

log()   { printf "\033[1;36m▶\033[0m %s\n" "$*"; }
warn()  { printf "\033[1;33m⚠\033[0m %s\n" "$*"; }
err()   { printf "\033[1;31m✗\033[0m %s\n" "$*" >&2; }
ok()    { printf "\033[1;32m✓\033[0m %s\n" "$*"; }

OS="$(uname -s)"
ARCH="$(uname -m)"

log "EthoVision worker installer"
log "OS: $OS · arch: $ARCH · env: $ENV_NAME"

# ---------- 1. Conda ----------
if command -v conda >/dev/null 2>&1; then
  ok "conda already installed — using it"
  eval "$(conda shell.bash hook)"
else
  log "conda not found — installing Miniconda to ~/miniconda3"
  TMPDIR="$(mktemp -d)"
  case "$OS" in
    Darwin)
      if [ "$ARCH" = "arm64" ]; then
        URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh"
      else
        URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh"
      fi
      ;;
    Linux)
      URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${ARCH}.sh"
      ;;
    *)
      err "Unsupported OS: $OS"; exit 1 ;;
  esac
  curl -fsSL "$URL" -o "$TMPDIR/miniconda.sh"
  bash "$TMPDIR/miniconda.sh" -b -p "$HOME/miniconda3"
  rm -rf "$TMPDIR"
  eval "$("$HOME/miniconda3/bin/conda" shell.bash hook)"
  ok "Miniconda installed"
fi

# ---------- 2. Env ----------
if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
  ok "env '$ENV_NAME' already exists — reusing"
else
  # Accept Anaconda's Terms of Service for default channels (added mid-2024).
  # Silently no-op on older conda that doesn't have the tos plugin.
  log "accepting Anaconda Terms of Service for default channels"
  for ch in \
    "https://repo.anaconda.com/pkgs/main" \
    "https://repo.anaconda.com/pkgs/r" \
    "https://repo.anaconda.com/pkgs/msys2"
  do
    conda tos accept --override-channels --channel "$ch" >/dev/null 2>&1 || true
  done

  log "creating env '$ENV_NAME' with Python $PY_VER"
  conda create -n "$ENV_NAME" "python=$PY_VER" -y
fi
conda activate "$ENV_NAME"
python -m pip install --upgrade pip >/dev/null

# ---------- 3. PyTorch ----------
TORCH_VARIANT=""
if [ "$OS" = "Darwin" ]; then
  TORCH_VARIANT="mac"
elif command -v nvidia-smi >/dev/null 2>&1; then
  DRIVER_MAJOR="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1 | awk -F. '{print $1}')"
  log "NVIDIA driver detected: $DRIVER_MAJOR"
  # Driver→CUDA wheel mapping (conservative)
  if   [ "$DRIVER_MAJOR" -ge 550 ]; then TORCH_VARIANT="cu124"
  elif [ "$DRIVER_MAJOR" -ge 525 ]; then TORCH_VARIANT="cu121"
  elif [ "$DRIVER_MAJOR" -ge 470 ]; then TORCH_VARIANT="cu118"
  else
    warn "NVIDIA driver too old for current PyTorch CUDA builds (need 470+). Falling back to CPU."
    TORCH_VARIANT="cpu"
  fi
else
  warn "no NVIDIA GPU detected — installing CPU-only PyTorch (tracking will be ~20× slower)"
  TORCH_VARIANT="cpu"
fi

log "installing PyTorch ($TORCH_VARIANT) — this is the big download (~2 GB)"
case "$TORCH_VARIANT" in
  mac)
    pip install torch torchvision ;;
  cpu)
    pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu ;;
  cu118|cu121|cu124)
    pip install torch torchvision --index-url "https://download.pytorch.org/whl/$TORCH_VARIANT" ;;
esac

# ---------- 4. DLC ----------
# [modelzoo] was removed in DLC 3; with that extra pip falls back to TF-based 2.x.
# DLC 3 is still release-candidate-only on PyPI → need --pre.
log "installing DeepLabCut 3.x (PyTorch-based, SuperAnimal built-in)"
pip install --pre "deeplabcut>=3.0.0rc1" requests

# ---------- 5. Worker script ----------
if [ ! -f "$HERE/ethovision_worker.py" ]; then
  if [ -n "${ETHOVISION_SERVER:-}" ]; then
    log "downloading ethovision_worker.py from $ETHOVISION_SERVER"
    curl -fsSL "$ETHOVISION_SERVER/worker/ethovision_worker.py" -o "$HERE/ethovision_worker.py"
  else
    warn "ethovision_worker.py is not next to install.sh, and ETHOVISION_SERVER is unset — skipping download"
  fi
fi

# ---------- 6. Launcher ----------
if [ -n "${ETHOVISION_SERVER:-}" ] && [ -n "${ETHOVISION_TOKEN:-}" ]; then
  log "writing run.sh with your server + token"
  cat > "$HERE/run.sh" <<RUNSH
#!/usr/bin/env bash
set -euo pipefail
HERE="\$(cd "\$(dirname "\${BASH_SOURCE[0]}")" && pwd)"
eval "\$(conda shell.bash hook)"
conda activate $ENV_NAME
export ETHOVISION_SERVER="$ETHOVISION_SERVER"
export ETHOVISION_TOKEN="$ETHOVISION_TOKEN"
exec python "\$HERE/ethovision_worker.py" "\$@"
RUNSH
  chmod +x "$HERE/run.sh"
  ok "launcher written: $HERE/run.sh"
fi

# ---------- 7. Sanity checks ----------
log "verifying install"
python - <<'PYEOF' || true
import sys
def check(name, mod):
    try:
        m = __import__(mod)
        print(f"  {name}: {getattr(m, '__version__', '?')}")
    except Exception as e:
        print(f"  {name}: FAILED ({e})")
check("torch", "torch")
check("deeplabcut", "deeplabcut")
check("requests", "requests")
try:
    import torch
    print(f"  CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"  GPU: {torch.cuda.get_device_name(0)}")
except Exception:
    pass
PYEOF

ok "Install complete."
echo
if [ -f "$HERE/run.sh" ]; then
  echo "Start the worker with:"
  echo "    bash $HERE/run.sh"
else
  echo "Start the worker with:"
  echo "    conda activate $ENV_NAME"
  echo "    export ETHOVISION_SERVER=http://<your-server>:3012"
  echo "    export ETHOVISION_TOKEN=<your-token>"
  echo "    python $HERE/ethovision_worker.py"
fi