#!/usr/bin/env bash
# install.sh — write a worker .env from the public gateway's bootstrap document.
#
# Usage:
#   ./install.sh --wallet <SOLANA_PUBKEY> [--gpu --compute-cap 89]
#                [--gateway-url https://api.autarcho.com] [--force]
#                [--accept-eula]
#
# This script only writes .env. It does not download images or start containers.

set -euo pipefail

VERSION="0.1.0"
REGISTRY="ghcr.io/czeti"
EULA_VERSION="1.0"
EULA_URL="https://autarcho.com/workers/EULA.md"
# SHA-256 of worker-docs/EULA.md for EULA_VERSION. Update both together.
EULA_SHA256="e1c81bcfa3b550d518c62dfeacebbb19d278347334f6858f6d8a40c91c37ed3a"
DEFAULT_GATEWAY="https://api.autarcho.com"
# Pinned out-of-band in README.md and worker.env.example. A bootstrap response
# for the default gateway must name this PeerId, not one supplied by the network.
PUBLIC_GATEWAY_PEER_ID="12D3KooWPEc9EMvcNttN1G6CSkNQm4gLKT9LEE5uKBJrCS2Wz4nE"
GATEWAY_URL="${AUTARCHO_GATEWAY_URL:-$DEFAULT_GATEWAY}"
EXPECTED_PEER_ID="${AUTARCHO_EXPECTED_GATEWAY_PEER_ID:-}"
WALLET="${WORKER_SOLANA_WALLET:-}"
GPU=0
COMPUTE_CAP="${WORKER_CUDA_COMPUTE_CAP:-}"
FORCE=0
DRY_RUN=0
ACCEPT_EULA=0

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

usage() {
  cat <<'EOF'
Usage: ./install.sh --wallet <SOLANA_PUBKEY> [options]

Required:
  --wallet <PUBKEY>       Solana mainnet USDC payout address

Options:
  --gpu                   Configure the CUDA worker override
  --compute-cap <N>       GPU compute capability: 75, 80, 86, 89, or 90
                          (required with --gpu)
  --gateway-url <URL>     Gateway base URL (default: https://api.autarcho.com)
  --expect-peer-id <ID>   Require this gateway PeerId (defaults to the pinned
                          public gateway for the default URL; mandatory for a
                          custom --gateway-url)
  --accept-eula           Confirm that you have read and accept the Autarcho
                          Worker EULA. Required for non-interactive installs.
  --force                 Overwrite an existing .env
  --dry-run               Validate and print the resolved values; write nothing
  -h, --help              Show this help

The script displays the Autarcho Worker EULA and records acceptance in .env.
It fetches <gateway-url>/bootstrap.json, pins the gateway PeerId and libp2p
address, and writes .env next to this script. It needs curl and jq.
EOF
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --wallet)
      [ "$#" -ge 2 ] || die "--wallet needs a value"
      WALLET="$2"
      shift 2
      ;;
    --gpu)
      GPU=1
      shift
      ;;
    --compute-cap)
      [ "$#" -ge 2 ] || die "--compute-cap needs a value"
      COMPUTE_CAP="$2"
      shift 2
      ;;
    --gateway-url)
      [ "$#" -ge 2 ] || die "--gateway-url needs a value"
      GATEWAY_URL="$2"
      shift 2
      ;;
    --expect-peer-id)
      [ "$#" -ge 2 ] || die "--expect-peer-id needs a value"
      EXPECTED_PEER_ID="$2"
      shift 2
      ;;
    --accept-eula)
      ACCEPT_EULA=1
      shift
      ;;
    --force)
      FORCE=1
      shift
      ;;
    --dry-run)
      DRY_RUN=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      die "unknown argument: $1 (use --help)"
      ;;
  esac
done

[ -n "$WALLET" ] || die "--wallet is required"
[[ "$WALLET" =~ ^[1-9A-HJ-NP-Za-km-z]{32,44}$ ]] || die "wallet is not a valid Solana address"

if [ "$GPU" -eq 1 ]; then
  case "$COMPUTE_CAP" in
    75|80|86|89|90) ;;
    *) die "--gpu requires --compute-cap 75, 80, 86, 89, or 90" ;;
  esac
else
  [ -z "$COMPUTE_CAP" ] || die "--compute-cap only applies with --gpu"
fi

GATEWAY_URL="${GATEWAY_URL%/}"
case "$GATEWAY_URL" in
  https://*) ;;
  *) die "gateway URL must start with https://" ;;
esac

if [ -z "$EXPECTED_PEER_ID" ]; then
  if [ "$GATEWAY_URL" = "$DEFAULT_GATEWAY" ]; then
    EXPECTED_PEER_ID="$PUBLIC_GATEWAY_PEER_ID"
  else
    die "custom gateway URL requires --expect-peer-id (or AUTARCHO_EXPECTED_GATEWAY_PEER_ID)"
  fi
fi

if [ "$DRY_RUN" -ne 1 ] && [ -e .env ] && [ "$FORCE" -ne 1 ]; then
  die ".env already exists; use --force to overwrite it"
fi

command -v curl >/dev/null 2>&1 || die "curl is required"
command -v jq >/dev/null 2>&1 || die "jq is required (install jq, then re-run)"

if [ "$DRY_RUN" -ne 1 ] && [ "$ACCEPT_EULA" -ne 1 ]; then
  printf '\nAutarcho Worker EULA %s\n%s\n\n' "$EULA_VERSION" "$EULA_URL"
  if [ ! -t 0 ]; then
    die "EULA acceptance required for non-interactive installs; re-run with --accept-eula"
  fi
  printf 'Type "accept" to agree: '
  if ! read -r REPLY; then
    die "EULA not accepted; no .env written"
  fi
  case "$REPLY" in
    accept|Accept|ACCEPT) ;;
    *) die "EULA not accepted; no .env written" ;;
  esac
fi

if [ "$DRY_RUN" -eq 1 ]; then
  EULA_ACCEPTED_AT=""
else
  EULA_ACCEPTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
fi

BOOTSTRAP="$(mktemp)"
trap 'rm -f "$BOOTSTRAP"' EXIT

printf 'Fetching %s/bootstrap.json\n' "$GATEWAY_URL"
curl -fsSL --max-time 20 "$GATEWAY_URL/bootstrap.json" -o "$BOOTSTRAP"

grep -q '"peer_id"' "$BOOTSTRAP" || die "bootstrap response has no peer_id"
PEER_ID="$(jq -er '.peer_id' "$BOOTSTRAP")"
ADDR="$(jq -er '.addrs[0]' "$BOOTSTRAP")"
ANNOUNCED_URL="$(jq -er '.gateway_url' "$BOOTSTRAP")"

[ -n "$PEER_ID" ] || die "bootstrap response has an empty peer_id"
[ -n "$ADDR" ] || die "bootstrap response has no dialable addrs"
[ -n "$ANNOUNCED_URL" ] || die "bootstrap response has no gateway_url"

if [ "$PEER_ID" != "$EXPECTED_PEER_ID" ]; then
  die "gateway PeerId mismatch: expected $EXPECTED_PEER_ID, got $PEER_ID"
fi

if [ "$ANNOUNCED_URL" != "$GATEWAY_URL" ]; then
  die "bootstrap gateway_url '$ANNOUNCED_URL' does not match '$GATEWAY_URL'"
fi

if [[ "$ADDR" == */p2p/* ]]; then
  ADDR_PEER_ID="${ADDR##*/p2p/}"
  [ "$ADDR_PEER_ID" = "$PEER_ID" ] || die "bootstrap address names PeerId '$ADDR_PEER_ID', expected '$PEER_ID'"
else
  ADDR="${ADDR%/}/p2p/$PEER_ID"
fi

if [ "$GPU" -eq 1 ]; then
  GPU_DEVICE="cuda"
else
  GPU_DEVICE="cpu"
fi

MAX_MEMORY="${WORKER_MAX_MEMORY_BYTES:-12884901888}"
MEM_LIMIT="${WORKER_MEM_LIMIT:-16g}"
MAX_CONCURRENT="${WORKER_MAX_CONCURRENT_INFERENCE:-1}"
MAX_GPU_CACHE="${WORKER_MAX_GPU_CACHE_BYTES:-2000000000}"
WORKER_DATA_DIR="${WORKER_DATA:-worker-data}"

if [ "$DRY_RUN" -eq 1 ]; then
  printf '\nDry run: no file written.\n'
  printf '  gateway URL:  %s\n' "$GATEWAY_URL"
  printf '  gateway peer: %s\n' "$PEER_ID"
  printf '  p2p address:  %s\n' "$ADDR"
  printf '  wallet:       %s\n' "$WALLET"
  if [ "$ACCEPT_EULA" -eq 1 ]; then
    printf '  EULA:         version %s, sha256 %s (not recorded in dry run)\n' "$EULA_VERSION" "$EULA_SHA256"
  else
    printf '  EULA:         version %s, sha256 %s (acceptance required before writing .env)\n' "$EULA_VERSION" "$EULA_SHA256"
  fi
  printf '  image:        %s/node:%s' "$REGISTRY" "$VERSION"
  if [ "$GPU" -eq 1 ]; then
    printf '%s' "-cuda$COMPUTE_CAP"
  fi
  printf '\n'
  exit 0
fi

umask 077
cat > .env <<EOF
# Generated by Autarcho install.sh. Keep this file private to the host.
TAG=$VERSION
REGISTRY=$REGISTRY

# Autarcho Worker EULA acceptance (recorded at install time).
EULA_VERSION=$EULA_VERSION
EULA_URL=$EULA_URL
EULA_SHA256=$EULA_SHA256
EULA_ACCEPTED_AT=$EULA_ACCEPTED_AT

GATEWAY_URL=$GATEWAY_URL
GATEWAY_P2P_ADDR=$ADDR
GATEWAY_PEER_ID=$PEER_ID

WORKER_SOLANA_WALLET=$WALLET

WORKER_MAX_MEMORY_BYTES=$MAX_MEMORY
WORKER_MEM_LIMIT=$MEM_LIMIT
WORKER_MAX_CONCURRENT_INFERENCE=$MAX_CONCURRENT
WORKER_GPU_DEVICE=$GPU_DEVICE
WORKER_CUDA_COMPUTE_CAP=$COMPUTE_CAP
WORKER_MAX_GPU_CACHE_BYTES=$MAX_GPU_CACHE
WORKER_DATA=$WORKER_DATA_DIR

RUST_LOG=info
EOF

printf '\nWrote .env for %s\n' "$GATEWAY_URL"
printf '  gateway peer: %s\n' "$PEER_ID"
printf '  p2p address:  %s\n' "$ADDR"
printf '  wallet:       %s\n' "$WALLET"
printf '  EULA:         version %s, sha256 %s, accepted at %s\n' "$EULA_VERSION" "$EULA_SHA256" "$EULA_ACCEPTED_AT"
printf '  image:        %s/node:%s' "$REGISTRY" "$VERSION"
if [ "$GPU" -eq 1 ]; then
  printf '%s' "-cuda$COMPUTE_CAP"
fi
printf '\n\nNext:\n'
printf '  docker compose pull\n'
if [ "$GPU" -eq 1 ]; then
  printf '  docker compose -f docker-compose.yml -f docker-compose.gpu.yml up -d\n'
else
  printf '  docker compose up -d\n'
fi
printf '  docker compose logs -f node\n\n'
printf 'Data: %s holds your worker identity. Back it up; set WORKER_DATA to an\n' "$WORKER_DATA_DIR"
printf 'absolute host path in .env if you want a simple bind-mount backup.\n'