#!/bin/sh
set -eu

usage() {
  cat <<'EOF'
Usage:
  curl -fsSL https://get.antgpu.ai | sh -s -- <ENROLLMENT_TOKEN>

Environment overrides:
  ANTGPU_CONNECT_MANIFEST_URL        (default: https://api.antgpu.ai/internal/connect/connector/manifest)
  ANTGPU_CONNECT_WORKDIR             (default Linux: /var/lib/antgpu/connect, non-Linux: $HOME/.antgpu/connect)
  ANTGPU_CONNECT_SERVICE_NAME        (default: antgpu-connect)
  ANTGPU_CONNECT_PREREQ_ASSUME_YES   (default: false; continue install when host checks warn/fail)
  ANTGPU_CONNECT_SHOW_LOGS           (default: prompt; true/false to skip prompt)
EOF
}

log() {
  printf '%s\n' "$*"
}

warn() {
  printf 'warning: %s\n' "$*" >&2
}

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

need_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "required command not found: $1"
}

is_root() {
  [ "$(id -u)" -eq 0 ]
}

as_root() {
  if is_root; then
    "$@"
    return 0
  fi
  if command -v sudo >/dev/null 2>&1; then
    sudo "$@"
    return 0
  fi
  return 127
}

sha256_file() {
  f="$1"
  if command -v sha256sum >/dev/null 2>&1; then
    sha256sum "$f" | awk '{print tolower($1)}'
    return 0
  fi
  if command -v shasum >/dev/null 2>&1; then
    shasum -a 256 "$f" | awk '{print tolower($1)}'
    return 0
  fi
  if command -v openssl >/dev/null 2>&1; then
    openssl dgst -sha256 "$f" | awk -F'= ' '{print tolower($2)}'
    return 0
  fi
  return 1
}

if [ "${1-}" = "-h" ] || [ "${1-}" = "--help" ]; then
  usage
  exit 0
fi

TOKEN="${1-}"
[ -n "$TOKEN" ] || {
  usage
  die "enrollment token argument is required"
}

need_cmd uname
need_cmd curl
need_cmd awk
need_cmd mktemp
need_cmd install
need_cmd tr

PYTHON_BIN=""
if command -v python3 >/dev/null 2>&1; then
  PYTHON_BIN="python3"
elif command -v python >/dev/null 2>&1; then
  PYTHON_BIN="python"
else
  die "python3 (or python) is required for manifest parsing"
fi

OS_RAW="$(uname -s)"
ARCH_RAW="$(uname -m)"

case "$OS_RAW" in
  Linux) OS="linux" ;;
  Darwin) OS="darwin" ;;
  *) die "unsupported OS: $OS_RAW" ;;
esac

case "$ARCH_RAW" in
  x86_64|amd64) ARCH="amd64" ;;
  aarch64|arm64) ARCH="arm64" ;;
  *) die "unsupported architecture: $ARCH_RAW" ;;
esac

if [ "$OS" != "linux" ] || [ "$ARCH" != "amd64" ]; then
  warn "primary supported platform is linux/amd64 (NVIDIA hosts). Current: ${OS}/${ARCH}."
fi

PREREQ_ROWS=""
PREREQ_HAS_ISSUE=0

truthy() {
  case "$(printf '%s' "${1-}" | tr '[:upper:]' '[:lower:]')" in
    1|true|yes|y|on) return 0 ;;
    *) return 1 ;;
  esac
}

add_prereq() {
  status="$1"
  name="$2"
  details="$3"
  PREREQ_ROWS="${PREREQ_ROWS}${status}	${name}	${details}
"
  case "$status" in
    warn|fail) PREREQ_HAS_ISSUE=1 ;;
  esac
}

print_prereq_summary() {
  log ""
  log "Host prerequisite summary:"
  printf '%-8s %-26s %s\n' "STATUS" "CHECK" "DETAILS"
  printf '%-8s %-26s %s\n' "------" "-----" "-------"
  printf '%s' "$PREREQ_ROWS" | while IFS='	' read -r status name details; do
    [ -n "$status" ] || continue
    printf '%-8s %-26s %s\n' "$status" "$name" "$details"
  done
  log ""
}

confirm_prereq_continue() {
  [ "$PREREQ_HAS_ISSUE" -eq 1 ] || return 0
  if truthy "${ANTGPU_CONNECT_PREREQ_ASSUME_YES:-}"; then
    warn "continuing despite prerequisite warnings/failures because ANTGPU_CONNECT_PREREQ_ASSUME_YES is set"
    return 0
  fi
  if [ ! -r /dev/tty ]; then
    die "host prerequisites are not fully satisfied; set ANTGPU_CONNECT_PREREQ_ASSUME_YES=true to continue anyway in non-interactive mode"
  fi
  printf 'Continue anyway? [y/N] ' >/dev/tty
  read -r answer </dev/tty || answer=""
  case "$answer" in
    y|Y|yes|YES|Yes) return 0 ;;
    *) die "aborted before enrollment; token was not used by this installer" ;;
  esac
}

should_show_logs() {
  value="${ANTGPU_CONNECT_SHOW_LOGS:-}"
  if [ -n "$value" ]; then
    truthy "$value"
    return $?
  fi
  [ -r /dev/tty ] || return 1
  printf 'Show live connector logs now? [Y/n] ' >/dev/tty
  read -r answer </dev/tty || answer=""
  case "$answer" in
    n|N|no|NO|No) return 1 ;;
    *) return 0 ;;
  esac
}

show_systemd_logs_prompt() {
  service_name="$1"
  log "Recent service status:"
  as_root systemctl --no-pager --lines=20 status "${service_name}.service" || true
  log ""
  log "Logs command: journalctl -u ${service_name} -f"
  if should_show_logs; then
    as_root journalctl -u "${service_name}" -f
  fi
}

check_host_prerequisites() {
  log "Checking host prerequisites..."

  docker_cli_ok=0
  docker_daemon_ok=0
  if command -v docker >/dev/null 2>&1; then
    docker_cli_ok=1
    add_prereq ok "Docker CLI" "docker command found"
  else
    add_prereq fail "Docker CLI" "Docker is required but was not found in PATH"
  fi

  if [ "$docker_cli_ok" -eq 1 ]; then
    if docker info >/dev/null 2>&1; then
      docker_daemon_ok=1
      add_prereq ok "Docker daemon" "daemon/socket reachable"
    else
      add_prereq fail "Docker daemon" "daemon/socket not reachable; start Docker and make sure this user can access the socket"
    fi
  else
    add_prereq fail "Docker daemon" "not checked because Docker CLI is missing"
  fi

  if [ "$OS" = "linux" ]; then
    if command -v nvidia-smi >/dev/null 2>&1; then
      add_prereq ok "NVIDIA driver" "nvidia-smi found"
    else
      add_prereq warn "NVIDIA driver" "nvidia-smi not found; GPU workloads and benchmarks may fail"
    fi

    gpu_probe_image="ubuntu:22.04"
    if [ "$docker_daemon_ok" -eq 1 ]; then
      if ! docker image inspect "$gpu_probe_image" >/dev/null 2>&1; then
        log "Docker image $gpu_probe_image not found locally; Docker will download it for GPU runtime check."
      fi
      if command -v timeout >/dev/null 2>&1; then
        NVIDIA_DOCKER_CMD="timeout 120 docker run --rm --gpus all $gpu_probe_image nvidia-smi"
      else
        NVIDIA_DOCKER_CMD="docker run --rm --gpus all $gpu_probe_image nvidia-smi"
      fi
      if sh -c "$NVIDIA_DOCKER_CMD" >/dev/null 2>&1; then
        add_prereq ok "Docker GPU runtime" "Ubuntu GPU container ran nvidia-smi with --gpus all"
      else
        add_prereq fail "Docker GPU runtime" "GPU container probe failed; install/configure NVIDIA Container Toolkit or verify Docker can pull $gpu_probe_image"
      fi
    else
      add_prereq fail "Docker GPU runtime" "not checked because Docker daemon is unavailable"
    fi
  fi

  print_prereq_summary
  if [ "$PREREQ_HAS_ISSUE" -eq 1 ]; then
    warn "Some prerequisites are not fully working. GPU rentals/benchmarks may fail until fixed."
    warn "NVIDIA Container Toolkit docs: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html"
  fi
  warn "If enrollment succeeds but tunnel never connects, VPN/proxy/firewall filtering may block UDP overlay tunnels. Disable VPN or try another network/VPN endpoint."
  confirm_prereq_continue
}

check_host_prerequisites

MANIFEST_URL="${ANTGPU_CONNECT_MANIFEST_URL:-https://api.antgpu.ai/internal/connect/connector/manifest}"
SERVICE_NAME="${ANTGPU_CONNECT_SERVICE_NAME:-antgpu-connect}"

if [ "$OS" = "linux" ]; then
  DEFAULT_WORKDIR="/var/lib/antgpu/connect"
else
  DEFAULT_WORKDIR="${HOME}/.antgpu/connect"
fi
WORKDIR="${ANTGPU_CONNECT_WORKDIR:-$DEFAULT_WORKDIR}"

TMP_DIR="$(mktemp -d)"
trap 'rm -rf "$TMP_DIR"' EXIT INT TERM

MANIFEST_PATH="$TMP_DIR/manifest.json"
log "Fetching connector manifest..."
log "Manifest URL: $MANIFEST_URL"
curl -fsSL "$MANIFEST_URL" -o "$MANIFEST_PATH"

MANIFEST_ROW="$("$PYTHON_BIN" - "$OS" "$ARCH" "$MANIFEST_PATH" <<'PY'
import json, os, sys
os_name = sys.argv[1]
arch = sys.argv[2]
manifest_path = sys.argv[3]

with open(manifest_path, "r", encoding="utf-8") as f:
    m = json.load(f)

version = (m.get("version") or "").strip()
if not version:
    raise SystemExit("manifest has empty version")

for item in m.get("binaries", []):
    if (item.get("os") or "").strip() == os_name and (item.get("arch") or "").strip() == arch:
        url = (item.get("url") or "").strip()
        sha = (item.get("sha256") or "").strip().lower()
        if not url or not sha:
            raise SystemExit("manifest entry missing url/sha256")
        filename = os.path.basename(url)
        if not filename:
            raise SystemExit("manifest URL has no filename")
        print("\t".join([version, url, sha, filename]))
        raise SystemExit(0)

raise SystemExit(f"no manifest entry for {os_name}/{arch}")
PY
)"

VERSION="$(printf '%s' "$MANIFEST_ROW" | awk -F '\t' '{print $1}')"
BIN_URL="$(printf '%s' "$MANIFEST_ROW" | awk -F '\t' '{print $2}')"
EXPECTED_SHA="$(printf '%s' "$MANIFEST_ROW" | awk -F '\t' '{print $3}')"

download_file() {
  url="$1"
  out="$2"
  if [ -t 2 ]; then
    curl -fL --progress-bar "$url" -o "$out"
  else
    curl -fsSL "$url" -o "$out"
  fi
}

BIN_TMP="$TMP_DIR/connect"
log "Downloading connector ${VERSION} for ${OS}/${ARCH}..."
log "Binary URL: $BIN_URL"
download_file "$BIN_URL" "$BIN_TMP"
log "Verifying connector checksum..."

ACTUAL_SHA="$(sha256_file "$BIN_TMP" || true)"
[ -n "$ACTUAL_SHA" ] || die "unable to compute sha256 (need sha256sum, shasum, or openssl)"
[ "$ACTUAL_SHA" = "$EXPECTED_SHA" ] || die "sha256 mismatch for connector binary"

INSTALL_PATH=""
if [ "$OS" = "linux" ]; then
  if as_root true >/dev/null 2>&1; then
    INSTALL_PATH="/usr/local/bin/connect"
    as_root install -d -m 0755 /usr/local/bin
    as_root install -m 0755 "$BIN_TMP" "$INSTALL_PATH"
  else
    INSTALL_PATH="${HOME}/.local/bin/connect"
    install -d -m 0755 "${HOME}/.local/bin"
    install -m 0755 "$BIN_TMP" "$INSTALL_PATH"
    warn "sudo not available; installed to ${INSTALL_PATH} and skipping systemd setup"
  fi
else
  INSTALL_PATH="${HOME}/.local/bin/connect"
  install -d -m 0755 "${HOME}/.local/bin"
  install -m 0755 "$BIN_TMP" "$INSTALL_PATH"
fi

INSTALLED_VERSION="$($INSTALL_PATH --version 2>/dev/null | awk 'NR==1 {print $1}')"
[ "$INSTALLED_VERSION" = "$VERSION" ] || die "installed connector version mismatch: manifest=${VERSION} binary=${INSTALLED_VERSION:-unknown}"

if [ "$OS" = "linux" ] && command -v systemctl >/dev/null 2>&1 && as_root true >/dev/null 2>&1; then
  as_root systemctl stop "${SERVICE_NAME}.service" >/dev/null 2>&1 || true
  as_root rm -rf "$WORKDIR"
  as_root install -d -m 0700 "$WORKDIR"
  as_root install -d -m 0755 /var/log/antgpu

  ENV_FILE="/etc/default/${SERVICE_NAME}"
  UNIT_FILE="/etc/systemd/system/${SERVICE_NAME}.service"
  LOG_PATH="/var/log/antgpu/connect.log"

  {
    printf 'ANTGPU_CONNECT_TOKEN=%s\n' "$TOKEN"
    printf 'CONNECT_LOG_PATH=%s\n' "$LOG_PATH"
  } | as_root tee "$ENV_FILE" >/dev/null
  as_root chmod 0600 "$ENV_FILE"

  cat <<EOF | as_root tee "$UNIT_FILE" >/dev/null
[Unit]
Description=Antgpu Connect
After=network-online.target docker.service
Wants=network-online.target

[Service]
Type=simple
WorkingDirectory=$WORKDIR
EnvironmentFile=-$ENV_FILE
ExecStart=/bin/sh -lc 'if [ -f .auth ]; then exec $INSTALL_PATH --foreground; elif [ -n "\${ANTGPU_CONNECT_TOKEN:-}" ]; then exec $INSTALL_PATH --foreground "\$ANTGPU_CONNECT_TOKEN"; else echo "missing .auth and ANTGPU_CONNECT_TOKEN"; exit 1; fi'
Restart=on-failure
RestartSec=3
TimeoutStopSec=20

[Install]
WantedBy=multi-user.target
EOF

  as_root systemctl daemon-reload
  as_root systemctl enable --now "${SERVICE_NAME}.service"

  if as_root systemctl is-active --quiet "${SERVICE_NAME}.service"; then
    log "connector ${VERSION} installed at ${INSTALL_PATH}"
    log "service ${SERVICE_NAME} is active (systemd)"
    show_systemd_logs_prompt "$SERVICE_NAME"
  else
    warn "service ${SERVICE_NAME} did not become active; check logs:"
    warn "journalctl -u ${SERVICE_NAME} -n 100 --no-pager"
    exit 1
  fi
else
  rm -rf "$WORKDIR"
  install -d -m 0700 "$WORKDIR"
  (
    cd "$WORKDIR"
    CONNECT_LOG_PATH="$WORKDIR/connect.log" "$INSTALL_PATH" "$TOKEN"
  )
  log "connector ${VERSION} installed at ${INSTALL_PATH}"
  log "started in built-in daemon mode (no systemd configured)"
fi