#!/usr/bin/env bash
#
# DeepLM Ops installer for macOS and Ubuntu.
#
#   curl -fsSL https://deeplm.ai/dlm-ops-dev/install.sh | bash
#   curl -fsSL https://deeplm.ai/dlm-ops-dev/install.sh | bash -s -- --my-machine
#
# Safe to re-run. Each step checks its own state and skips work that is done,
# so this same command installs, upgrades, restarts and reopens DeepLM Ops.
#
# Everything lives in functions and main runs on the last line, so a download
# that is cut off part way does nothing.

set -Eeuo pipefail

VERSION="${DLM_OPS_VERSION:-0.2.0-rc.4}"
IMAGE="${DLM_OPS_IMAGE:-ghcr.io/deeplm/dlm-ops-dev:$VERSION}"
CONTAINER="dlm-ops"
PORT="8090"
STATE_DIR="$HOME/.dlm-ops"
MIN_CPU_MEMORY_GB=8
MIN_VRAM_MB=4608
DOCKER_START_TIMEOUT=300
READY_TIMEOUT="${DLM_OPS_READY_TIMEOUT:-1800}"
INSTALL_PAGE="https://deeplm.ai/dlm-ops-dev"

# Docker Desktop minimum version (macOS only).
MIN_DESKTOP_VERSION="4.71.0"

MODE="demo"
OPEN_BROWSER=1
ENGINE="cpu"
OS=""  # "macos" or "linux", set by detect_os

# Docker Desktop's CLI may not be on PATH yet, especially right after install.
PATH="$PATH:/opt/homebrew/bin:/usr/local/bin:$HOME/.docker/bin:/Applications/Docker.app/Contents/Resources/bin"

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

usage() {
  cat <<EOF
Usage: install.sh [--demo | --my-machine] [--no-open]

  --demo         Include the recorded sample cluster (default)
  --my-machine   Use only this machine, with no demo cluster
  --no-open      Print the sign-in link instead of opening a browser
EOF
}

parse_args() {
  while [ "$#" -gt 0 ]; do
    case "$1" in
      --demo) MODE="demo" ;;
      --my-machine) MODE="my-machine" ;;
      --no-open) OPEN_BROWSER=0 ;;
      -h | --help)
        usage
        exit 0
        ;;
      *)
        usage >&2
        die "Unknown option: $1"
        ;;
    esac
    shift
  done
}

# True when version $1 is at least version $2.
version_at_least() {
  awk -v have="$1" -v want="$2" 'BEGIN {
    n = split(have, h, ".")
    m = split(want, w, ".")
    max = (n > m) ? n : m
    for (i = 1; i <= max; i++) {
      if (h[i] + 0 > w[i] + 0) exit 0
      if (h[i] + 0 < w[i] + 0) exit 1
    }
    exit 0
  }'
}

detect_os() {
  case "$(uname -s)" in
    Darwin) OS="macos" ;;
    Linux)
      if [ -f /etc/os-release ]; then
        # Accept Ubuntu, Debian, Pop!_OS, and other apt-based distros.
        if grep -qiE 'ubuntu|debian|pop|mint' /etc/os-release 2>/dev/null; then
          OS="linux"
        else
          OS="linux"
          say "Warning: this installer is tested on Ubuntu. Your distro may need manual adjustments."
        fi
      else
        die "Cannot detect your Linux distribution. This installer supports Ubuntu 20.04+."
      fi
      ;;
    *) die "This installer supports macOS and Ubuntu. For other platforms, see $INSTALL_PAGE" ;;
  esac
}

open_url() {
  if [ "$OS" = "macos" ]; then
    open "$1" 2>/dev/null || true
  else
    xdg-open "$1" 2>/dev/null || true
  fi
}

# ── Docker installation ─────────────────────────────────────────────────────

install_docker_macos() {
  if command -v brew >/dev/null 2>&1; then
    say "Installing Docker Desktop with Homebrew. macOS may ask for your password."
    brew install --cask docker-desktop
  else
    die "Docker Desktop is required. Install it from https://www.docker.com/products/docker-desktop/ and run this again."
  fi
}

install_docker_linux() {
  say "Installing Docker Engine. This may ask for your sudo password."

  if ! command -v curl >/dev/null 2>&1; then
    sudo apt-get update -qq && sudo apt-get install -y -qq curl >/dev/null
  fi

  # Use Docker's official convenience script.
  curl -fsSL https://get.docker.com | sudo sh

  # Let the current user run docker without sudo.
  sudo usermod -aG docker "$USER" 2>/dev/null || true
  say "Added $USER to the docker group."
  say "If 'docker' commands fail after install, log out and back in (or run: newgrp docker)."
}

start_docker_macos() {
  docker info >/dev/null 2>&1 && return 0

  say "Starting Docker Desktop. On first launch, accept its terms in the Docker window."
  open -a Docker 2>/dev/null ||
    die "Could not open Docker Desktop. Open it manually and run this again."

  local waited=0
  until docker info >/dev/null 2>&1; do
    if [ "$waited" -ge "$DOCKER_START_TIMEOUT" ]; then
      die "Docker did not start within $((DOCKER_START_TIMEOUT / 60)) minutes. Finish any prompts in Docker Desktop and run this again."
    fi
    sleep 3
    waited=$((waited + 3))
  done
}

start_docker_linux() {
  docker info >/dev/null 2>&1 && return 0

  say "Starting Docker Engine."
  sudo systemctl start docker 2>/dev/null ||
    sudo service docker start 2>/dev/null ||
    die "Could not start Docker. Run 'sudo systemctl start docker' manually and try again."

  local waited=0
  until docker info >/dev/null 2>&1; do
    if [ "$waited" -ge "$DOCKER_START_TIMEOUT" ]; then
      die "Docker did not start within $((DOCKER_START_TIMEOUT / 60)) minutes."
    fi
    sleep 3
    waited=$((waited + 3))
  done
}

ensure_docker() {
  step "Checking Docker"

  if ! command -v docker >/dev/null 2>&1; then
    if [ "$OS" = "macos" ] && [ ! -d /Applications/Docker.app ]; then
      install_docker_macos
    elif [ "$OS" = "linux" ]; then
      install_docker_linux
    else
      die "Docker is not installed. Install it from https://www.docker.com/ and run this again."
    fi
  fi

  if [ "$OS" = "macos" ]; then
    start_docker_macos
  else
    start_docker_linux
  fi

  choose_engine
}

# ── Engine selection ─────────────────────────────────────────────────────────

check_nvidia_gpu_linux() {
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    return 1
  fi
  nvidia-smi >/dev/null 2>&1 || return 1

  # Check that docker can use the NVIDIA runtime.
  if docker info 2>/dev/null | grep -qi 'nvidia'; then
    return 0
  fi
  # Also check for the runtime directly.
  if docker info --format '{{json .Runtimes}}' 2>/dev/null | grep -q '"nvidia"'; then
    return 0
  fi
  return 1
}

check_vram_linux() {
  local free_mb
  free_mb="$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ')"
  if [ -z "$free_mb" ]; then
    return 0  # can't check, proceed anyway
  fi
  if [ "$free_mb" -lt "$MIN_VRAM_MB" ]; then
    say ""
    say "Warning: only ${free_mb} MiB of GPU memory free (need ${MIN_VRAM_MB} MiB)."
    say "Other processes may be using your GPU. Check with: nvidia-smi"
    say "Common fix: stop Ollama (systemctl stop ollama) or other GPU processes."
    say ""
    say "Continuing with CPU mode. Restart after freeing VRAM to use the GPU."
    return 1
  fi
  return 0
}

choose_engine() {
  if [ "$OS" = "macos" ]; then
    choose_engine_macos
  else
    choose_engine_linux
  fi
}

choose_engine_macos() {
  local platform desktop_version memory_bytes memory_gb
  platform="$(docker version --format '{{.Server.Platform.Name}}' 2>/dev/null || true)"
  desktop_version="$(printf '%s' "$platform" |
    sed -n 's/^Docker Desktop \([0-9][0-9.]*\).*/\1/p')"

  if [ "$(uname -m)" = "arm64" ] && [ -n "$desktop_version" ]; then
    if ! version_at_least "$desktop_version" "$MIN_DESKTOP_VERSION"; then
      die "Docker Desktop $MIN_DESKTOP_VERSION or later is required, and you have $desktop_version. Update it with: brew upgrade --cask docker-desktop"
    fi
    ENGINE="dmr"
    say "Docker Desktop $desktop_version: using this Mac's GPU"
    return 0
  fi

  ENGINE="cpu"
  memory_bytes="$(docker info --format '{{.MemTotal}}' 2>/dev/null || echo 0)"
  memory_gb=$((memory_bytes / 1024 / 1024 / 1024))
  if [ $((memory_gb + 1)) -lt "$MIN_CPU_MEMORY_GB" ]; then
    die "The model will run on the CPU inside Docker, which needs ${MIN_CPU_MEMORY_GB} GB of Docker memory. Docker has about ${memory_gb} GB. Raise it in Docker Desktop under Settings > Resources, then run this again."
  fi
  say "Using the CPU inside Docker (${memory_gb} GB of memory available)"
}

choose_engine_linux() {
  if check_nvidia_gpu_linux && check_vram_linux; then
    ENGINE="nvidia"
    local gpu_name
    gpu_name="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)"
    say "NVIDIA GPU detected: ${gpu_name:-unknown}"
  else
    ENGINE="cpu"
    local memory_bytes memory_gb
    memory_bytes="$(docker info --format '{{.MemTotal}}' 2>/dev/null || echo 0)"
    memory_gb=$((memory_bytes / 1024 / 1024 / 1024))
    if [ $((memory_gb + 1)) -lt "$MIN_CPU_MEMORY_GB" ]; then
      die "CPU mode needs ${MIN_CPU_MEMORY_GB} GB of memory available to Docker. Docker reports about ${memory_gb} GB."
    fi
    say "Using CPU mode (${memory_gb} GB of memory available)"
    if ! command -v nvidia-smi >/dev/null 2>&1; then
      say "Tip: install NVIDIA drivers and the NVIDIA Container Toolkit for GPU acceleration."
    fi
  fi
}

# ── GitHub sign-in ───────────────────────────────────────────────────────────

# ── Image management ─────────────────────────────────────────────────────────

ensure_image() {
  step "Getting DeepLM Ops $VERSION"
  if docker image inspect "$IMAGE" >/dev/null 2>&1; then
    say "Already downloaded"
    return 0
  fi

  docker pull "$IMAGE" ||
    die "Could not pull $IMAGE. Check your internet connection and try again."
}

# ── Model runner (macOS Apple Silicon only) ──────────────────────────────────

ensure_model_runner() {
  [ "$ENGINE" = "dmr" ] || return 0
  step "Setting up the GPU runtime"
  local model_status
  model_status="$(docker model status 2>&1 || true)"
  if grep -qi 'is running' <<<"$model_status"; then
    say "Docker Model Runner is ready"
    return 0
  fi
  docker desktop enable model-runner
  say "Docker Model Runner enabled"
}

ensure_host_probe() {
  [ "$ENGINE" = "dmr" ] || return 0
  step "Setting up the host probe"

  local marker="$STATE_DIR/.probe-version"
  if [ -x "$STATE_DIR/bin/dlm-ops-probe" ] &&
    [ "$(cat "$marker" 2>/dev/null || true)" = "$VERSION" ]; then
    say "Already installed"
    return 0
  fi

  mkdir -p "$STATE_DIR/host"
  docker run --rm -v "$STATE_DIR:/out" "$IMAGE" install-host-probe
  "$STATE_DIR/bin/dlm-ops-probe" install
  printf '%s\n' "$VERSION" >"$marker"
  say "Installed"
}

# ── Container management ────────────────────────────────────────────────────

check_port_in_use() {
  if [ "$OS" = "macos" ]; then
    lsof -nP -iTCP:"$PORT" -sTCP:LISTEN >/dev/null 2>&1
  else
    ss -tlnp 2>/dev/null | grep -q ":$PORT " 2>/dev/null ||
      lsof -nP -iTCP:"$PORT" -sTCP:LISTEN >/dev/null 2>&1
  fi
}

start_new_container() {
  if check_port_in_use; then
    die "Port $PORT is already in use. Stop whatever is using it and run this again."
  fi

  local args=(
    run -d
    --name "$CONTAINER"
    -p "127.0.0.1:$PORT:$PORT"
    -v dlm-ops:/var/lib/dlm-ops
    --read-only
    --tmpfs /tmp
    --tmpfs "/run/dlm:uid=1000,gid=1000,mode=0700"
    --cap-drop ALL
    --pids-limit 256
    --restart unless-stopped
    --label "ai.deeplm.dlm-ops.config=$1"
  )

  case "$ENGINE" in
    dmr)
      args+=(-e DLM_OPS_ENGINE=dmr -v "$STATE_DIR/host:/host:ro")
      ;;
    nvidia)
      args+=(
        --runtime=nvidia
        -e NVIDIA_VISIBLE_DEVICES=all
        -e NVIDIA_DRIVER_CAPABILITIES=compute,utility
      )
      ;;
  esac

  args+=("$IMAGE")
  if [ "$MODE" = "demo" ]; then
    args+=(demo)
  fi

  docker "${args[@]}" >/dev/null
}

ensure_container() {
  step "Starting DeepLM Ops"

  local wanted="$VERSION/$MODE/$ENGINE" status current
  status="$(docker inspect -f '{{.State.Status}}' "$CONTAINER" 2>/dev/null || true)"
  if [ -n "$status" ]; then
    current="$(docker inspect \
      -f '{{index .Config.Labels "ai.deeplm.dlm-ops.config"}}' "$CONTAINER" \
      2>/dev/null || true)"
    if [ "$current" != "$wanted" ]; then
      say "Replacing the existing container. Your data is kept."
      docker rm -f "$CONTAINER" >/dev/null
      status=""
    fi
  fi

  case "$status" in
    running | restarting) say "Already running" ;;
    "")
      start_new_container "$wanted"
      say "Started"
      ;;
    *)
      docker start "$CONTAINER" >/dev/null
      say "Restarted"
      ;;
  esac
}

# ── Readiness ────────────────────────────────────────────────────────────────

container_logs() {
  docker logs --since "$1" "$CONTAINER" 2>&1 || true
}

hide_token() {
  sed 's/#token=[^[:space:]]*/#token=[hidden]/g'
}

wait_until_ready() {
  step "Waiting for the model"
  say "The first start downloads 2.55 GiB once. Later starts are much faster."

  local started_at waited=0 logs last_line printed_line="" status
  started_at="$(docker inspect -f '{{.State.StartedAt}}' "$CONTAINER")"

  while :; do
    logs="$(container_logs "$started_at")"
    SIGN_IN_URL="$(printf '%s\n' "$logs" |
      grep -Eo "http://127\.0\.0\.1:$PORT/#token=[^[:space:]\"']+" | tail -n 1 || true)"

    if [ -n "$SIGN_IN_URL" ] && grep -qi 'model ready' <<<"$logs"; then
      say "Model ready"
      return 0
    fi

    status="$(docker inspect -f '{{.State.Status}}' "$CONTAINER" 2>/dev/null || true)"
    if [ "$status" != "running" ] && [ "$status" != "restarting" ]; then
      printf '%s\n' "$logs" | tail -n 20 | hide_token >&2
      die "DeepLM Ops stopped during startup. Its last log lines are above."
    fi

    if [ "$waited" -ge "$READY_TIMEOUT" ]; then
      die "DeepLM Ops was not ready after $((READY_TIMEOUT / 60)) minutes. Run 'docker logs dlm-ops' to see why, and do not share any link it prints."
    fi

    last_line="$(printf '%s\n' "$logs" | tail -n 1 | hide_token)"
    if [ -n "$last_line" ] && [ "$last_line" != "$printed_line" ]; then
      say "$last_line"
      printed_line="$last_line"
    fi

    sleep 3
    waited=$((waited + 3))
  done
}

# ── Final output ─────────────────────────────────────────────────────────────

open_app() {
  step "DeepLM Ops is ready"
  if [ "$OPEN_BROWSER" -eq 1 ]; then
    open_url "$SIGN_IN_URL"
    say "Opened in your browser."
  else
    say "Open this link. It is a password, so do not paste it into chats, issues or screen shares:"
    say "$SIGN_IN_URL"
  fi

  printf '\n'
  say "To start DeepLM Ops again later:"
  say "  docker start dlm-ops && docker logs -f dlm-ops"
  say ""
  say "Or re-run this installer (it picks up where it left off)."

  printf '\n  Try asking:\n'
  if [ "$MODE" = "demo" ]; then
    say "  Why is my job 5102 still pending?"
    say "  Why did 5098 fail?"
    say "  Is anyone wasting a GPU on gpu03?"
  fi
  say "  Will Llama 3.1 8B fit on my GPU for inference with 8k context?"
}

# ── Main ─────────────────────────────────────────────────────────────────────

main() {
  parse_args "$@"
  detect_os
  ensure_docker
  ensure_image
  ensure_model_runner
  ensure_host_probe
  ensure_container
  wait_until_ready
  open_app
}

main "$@" </dev/null