#!/bin/sh
set -eu

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

Environment overrides:
  ANTHIVE_CONNECT_MANIFEST_URL   (default: https://api.anthive.ai/internal/connect/connector/manifest)
  ANTHIVE_CONNECT_WORKDIR        (default Linux: /var/lib/anthive/connect, non-Linux: $HOME/.anthive/connect)
  ANTHIVE_CONNECT_SERVICE_NAME   (default: anthive-connect)
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

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

if ! command -v docker >/dev/null 2>&1; then
  die "docker is required but not found in PATH"
fi

if [ "$OS" = "linux" ] && ! command -v nvidia-smi >/dev/null 2>&1; then
  warn "nvidia-smi not found; this host may not have NVIDIA runtime ready"
fi

MANIFEST_URL="${ANTHIVE_CONNECT_MANIFEST_URL:-https://api.anthive.ai/internal/connect/connector/manifest}"
SERVICE_NAME="${ANTHIVE_CONNECT_SERVICE_NAME:-anthive-connect}"

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

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

MANIFEST_PATH="$TMP_DIR/manifest.json"
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}')"

BIN_TMP="$TMP_DIR/connect"
curl -fsSL "$BIN_URL" -o "$BIN_TMP"

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

if [ "$OS" = "linux" ] && command -v systemctl >/dev/null 2>&1 && as_root true >/dev/null 2>&1; then
  as_root install -d -m 0755 "$WORKDIR"
  as_root install -d -m 0755 /var/log/anthive

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

  {
    printf 'ANTHIVE_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=Anthive 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 "\${ANTHIVE_CONNECT_TOKEN:-}" ]; then exec $INSTALL_PATH --foreground "\$ANTHIVE_CONNECT_TOKEN"; else echo "missing .auth and ANTHIVE_CONNECT_TOKEN"; exit 1; fi'
Restart=always
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)"
    log "logs: journalctl -u ${SERVICE_NAME} -f"
  else
    warn "service ${SERVICE_NAME} did not become active; check logs:"
    warn "journalctl -u ${SERVICE_NAME} -n 100 --no-pager"
    exit 1
  fi
else
  install -d -m 0755 "$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