#!/usr/bin/env bash
# Join a Linux NVIDIA GPU box to the Feya network as a provider node (Docker).
#
# Usage:
#   curl -sSL https://feya-ai.com/join.sh | FEYA_PAIRING_CODE=xxxx bash
#
# or from a local copy:
#   FEYA_PAIRING_CODE=xxxx bash join.sh
#   bash join.sh xxxx                 # code as first arg also works
#
# Where the code comes from: sign in at https://feya-ai.com (accounts are
# invite-gated during beta), click "Add device" — that mints a pairing code.
# Codes are SINGLE-USE and expire after 10 minutes, so mint one per box right
# before you run this. Running WITHOUT a code also works: the box still joins
# and serves inference, but its work receipts aren't attributed to your
# account until you pair it (re-run with a code; identity is preserved).
#
# What this script does (idempotent; safe to re-run):
#   1. Preflight: Linux x86_64, Docker daemon reachable, NVIDIA driver
#      (nvidia-smi), nvidia-container-toolkit wired into Docker.
#   2. Gets the provider image: tries `docker pull` first, falls back to
#      building locally from the public repo (see HONESTY note below).
#   3. Downloads the model GGUF (~4.9 GB, one-time, resumable).
#   4. Starts the node as a container: --restart unless-stopped, GPU
#      passthrough, outbound-only WSS tunnel to the coordinator (no
#      inbound ports, no port forwarding, NAT/firewall friendly).
#   5. Verifies the node actually shows up in the coordinator's provider
#      registry before declaring success.
#
# HONESTY NOTES (state of the world as of 2026-07-11):
#   - The provider image is published to ghcr.io by CI on every push to
#     main (.github/workflows/publish-image.yml). If the pull fails
#     anyway (registry outage, this box is offline from ghcr), the
#     script falls back to `docker build` from the repo, which takes
#     ~5-10 minutes per box (or once, if you build on one box and
#     `docker save | docker load` / push to your own registry).
#   - Model weight hashes are bootstrap-permissive during beta (the
#     approved-hashes registry has unpinned entries), so the download below
#     is size-checked but not yet hash-pinned by this script.
#
# Env knobs (all optional):
#   FEYA_PAIRING_CODE  pairing code from the dashboard (FEYA_INVITE and
#                      PAIRING_CODE accepted as aliases)
#   FEYA_HOME          host dir for models + node state   (default ~/.feya)
#   FEYA_IMAGE         image to pull                      (default ghcr.io/audrey-mcknight/pekito-standalone:latest)
#   FEYA_REPO          repo for the build fallback        (default https://github.com/audrey-mcknight/pekito.git)
#   FEYA_MODEL_URL     GGUF to serve (default Llama 3.1 8B Q4_K_M; boxes with
#                      <8 GB VRAM should point at a smaller model — see
#                      https://feya-ai.com/v1/models/catalog)
#   PEKITO_COORDINATOR_URL  coordinator (default https://feya-ai.com)
#   PEKITO_NODE_NAME        display name (default hostname)
set -euo pipefail

# ---- config knobs ---------------------------------------------------------
PAIRING_CODE="${1:-${FEYA_PAIRING_CODE:-${FEYA_INVITE:-${PAIRING_CODE:-}}}}"
COORD_URL="${PEKITO_COORDINATOR_URL:-https://feya-ai.com}"
FEYA_HOME="${FEYA_HOME:-$HOME/.feya}"
IMAGE="${FEYA_IMAGE:-ghcr.io/audrey-mcknight/pekito-standalone:latest}"
REPO_URL="${FEYA_REPO:-https://github.com/audrey-mcknight/pekito.git}"
NODE_NAME="${PEKITO_NODE_NAME:-$(hostname -s 2>/dev/null || echo feya-node)}"
CONTAINER_NAME="${FEYA_CONTAINER_NAME:-feya-node}"

MODEL_URL="${FEYA_MODEL_URL:-https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf}"
MODEL_FILE="$FEYA_HOME/models/$(basename "$MODEL_URL")"

log()  { printf "\033[1;32m[feya]\033[0m %s\n" "$*"; }
warn() { printf "\033[1;33m[feya]\033[0m %s\n" "$*" >&2; }
die()  { printf "\033[1;31m[feya]\033[0m %s\n" "$*" >&2; exit 1; }

# ---- preflight ------------------------------------------------------------
[[ "$(uname -s)" == "Linux" ]] || die "this script is Linux-only. Mac: install the Feya app (FeyaAI.dmg) from https://feya-ai.com/app/"
[[ "$(uname -m)" == "x86_64" ]] || die "x86_64 only (the CUDA image and prebuilt wheels are amd64)"

command -v curl >/dev/null 2>&1 || die "curl is required (apt-get install -y curl)"

command -v docker >/dev/null 2>&1 || die "Docker not found. Install it: https://docs.docker.com/engine/install/ubuntu/"
if ! docker info >/dev/null 2>&1; then
  die "Docker daemon unreachable. Is it running (systemctl start docker)? If you're not root, add yourself to the docker group or re-run with sudo."
fi

command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi not found — install the NVIDIA driver first (ubuntu-drivers autoinstall, then reboot)."
GPU_NAME="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || true)"
GPU_VRAM_MB="$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ' || echo 0)"
log "GPU: ${GPU_NAME:-unknown} (${GPU_VRAM_MB:-?} MB VRAM)"
if [[ "${GPU_VRAM_MB:-0}" =~ ^[0-9]+$ ]] && (( GPU_VRAM_MB < 8000 )) && [[ -z "${FEYA_MODEL_URL:-}" ]]; then
  warn "less than 8 GB VRAM — Llama 3.1 8B Q4_K_M may not fit fully on GPU."
  warn "consider: FEYA_MODEL_URL=https://huggingface.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF/resolve/main/Qwen2.5-1.5B-Instruct-Q4_K_M.gguf"
fi

# nvidia-container-toolkit: docker must know the nvidia runtime, or
# `--gpus all` fails at container start.
if ! docker info 2>/dev/null | grep -qi 'nvidia'; then
  warn "Docker does not list an nvidia runtime. Install nvidia-container-toolkit:"
  warn "  https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html"
  warn "  (Ubuntu: apt-get install -y nvidia-container-toolkit && nvidia-ctk runtime configure --runtime=docker && systemctl restart docker)"
  die "nvidia-container-toolkit missing — install it and re-run"
fi

if [[ -z "$PAIRING_CODE" && ! -f "$FEYA_HOME/state/claimed.json" ]]; then
  warn "no pairing code given — the box will join UNPAIRED. It serves inference"
  warn "but earnings aren't attributed to an account until you pair it."
  warn "Mint a code at https://feya-ai.com (Add device) and re-run with"
  warn "FEYA_PAIRING_CODE=<code> at any time; node identity is preserved."
fi

mkdir -p "$FEYA_HOME/models" "$FEYA_HOME/state"

# ---- image: pull, else build locally --------------------------------------
# HONEST: no image is published yet (see header). Pull is attempted first so
# this script keeps working unchanged the day images land on ghcr.io; today
# it falls through to the local build.
RUN_IMAGE=""
if docker image inspect "$IMAGE" >/dev/null 2>&1; then
  log "image $IMAGE already present locally"
  RUN_IMAGE="$IMAGE"
elif docker pull "$IMAGE" >/dev/null 2>&1; then
  log "pulled $IMAGE"
  RUN_IMAGE="$IMAGE"
else
  warn "could not pull $IMAGE (no published Feya image yet — expected during beta)"
  log "falling back to a local build from $REPO_URL (~5-10 min, one-time)"
  command -v git >/dev/null 2>&1 || die "git is required for the build fallback (apt-get install -y git)"
  SRC_DIR="$FEYA_HOME/src"
  # GIT_TERMINAL_PROMPT=0 + timeout: an unauthenticated clone/fetch of a
  # private $REPO_URL otherwise blocks forever on a credential prompt that
  # never arrives in this non-interactive script, indistinguishable from a
  # slow network. Fail fast and say why instead.
  if [[ -d "$SRC_DIR/.git" ]]; then
    log "updating source at $SRC_DIR"
    GIT_TERMINAL_PROMPT=0 timeout 60 git -C "$SRC_DIR" fetch --quiet origin main \
      || die "git fetch of $REPO_URL timed out or failed — if it's private, make sure this box has credentials cached (git credential-cache, or an SSH key on the remote)"
    git -C "$SRC_DIR" reset --quiet --hard origin/main
  else
    log "cloning $REPO_URL"
    rm -rf "$SRC_DIR"
    GIT_TERMINAL_PROMPT=0 timeout 60 git clone --quiet --depth 1 "$REPO_URL" "$SRC_DIR" \
      || die "git clone of $REPO_URL timed out or failed — if it's private, make sure this box has credentials cached (git credential-cache, or an SSH key on the remote)"
  fi
  RUN_IMAGE="feya-standalone:local"
  docker build -f "$SRC_DIR/docker/Dockerfile.standalone" -t "$RUN_IMAGE" "$SRC_DIR" \
    || die "docker build failed — see output above. Tip for a fleet: build once, then 'docker save' / push to your own registry."
  log "built $RUN_IMAGE"
fi

# GPU-in-Docker sanity check with the actual image (it's CUDA-based, so
# nvidia-smi is inside). Catches a mis-configured toolkit before we launch.
if ! docker run --rm --gpus all --entrypoint nvidia-smi "$RUN_IMAGE" -L >/dev/null 2>&1; then
  die "GPU passthrough test failed ('docker run --gpus all ... nvidia-smi'). Check nvidia-container-toolkit config (nvidia-ctk runtime configure --runtime=docker && systemctl restart docker) and re-run."
fi
log "GPU passthrough OK"

# ---- model download --------------------------------------------------------
# Size-validated: "file exists" isn't enough if an aborted run left a partial
# file. HEAD for Content-Length, compare, resume with curl -C - if short.
expected_size="$(curl -sfL -I "$MODEL_URL" 2>/dev/null | awk 'tolower($1) == "content-length:" {print $2}' | tr -d '\r' | tail -1)"
need_download=1
if [[ -f "$MODEL_FILE" ]]; then
  actual_size="$(stat -c%s "$MODEL_FILE" 2>/dev/null || echo 0)"
  if [[ -z "$expected_size" ]] || (( actual_size >= expected_size )); then
    log "model already present at $MODEL_FILE"
    need_download=0
  else
    warn "model is truncated ($actual_size / $expected_size bytes); resuming download"
  fi
fi
if (( need_download )); then
  log "downloading model (one-time, resumable): $(basename "$MODEL_FILE")"
  curl -L --retry 3 --retry-delay 2 -C - -o "$MODEL_FILE" "$MODEL_URL"
fi

# ---- launch -----------------------------------------------------------------
# State volume (/state) persists node_id, the Ed25519 receipt-signing key and
# the pairing marker across container recreations — losing these would change
# the node's durable identity and orphan its account binding.
# Tunnel backend is WSS: the node dials OUT to the coordinator over HTTPS and
# inference requests come back down that socket. No inbound ports published.
if docker ps -a --format '{{.Names}}' | grep -qx "$CONTAINER_NAME"; then
  log "removing existing container $CONTAINER_NAME (state survives in $FEYA_HOME)"
  docker rm -f "$CONTAINER_NAME" >/dev/null
fi

run_args=(
  -d --name "$CONTAINER_NAME"
  --restart unless-stopped
  --gpus all
  -v "$FEYA_HOME/models:/models:ro"
  -v "$FEYA_HOME/state:/state"
  -e PEKITO_HOME=/state
  -e PEKITO_COORDINATOR_URL="$COORD_URL"
  -e PEKITO_NODE_NAME="$NODE_NAME"
  -e PEKITO_MODEL_PATH="/models/$(basename "$MODEL_FILE")"
  -e PEKITO_TUNNEL_BACKEND=wss
)
if [[ -n "$PAIRING_CODE" ]]; then
  run_args+=(-e PEKITO_PAIRING_CODE="$PAIRING_CODE")
fi

log "starting node '$NODE_NAME' (container $CONTAINER_NAME)"
docker run "${run_args[@]}" "$RUN_IMAGE" >/dev/null

# ---- verify against the live network ----------------------------------------
# The node registers itself at $COORD_URL/v1/providers/register once the model
# is loaded; loading a 5 GB GGUF takes tens of seconds. Poll the public
# provider registry for our durable node_id.
log "waiting for the node to appear on the network (model load takes ~30-90s)"
deadline=$(( $(date +%s) + 180 ))
joined=0
NODE_ID=""
while (( $(date +%s) < deadline )); do
  # node_id is minted inside the container on first boot and lands in the
  # mounted state dir a moment after start — re-read until it exists.
  [[ -z "$NODE_ID" ]] && NODE_ID="$(cat "$FEYA_HOME/state/node_id" 2>/dev/null || true)"
  if ! docker ps --format '{{.Names}}' | grep -qx "$CONTAINER_NAME"; then
    warn "container exited — last log lines:"
    docker logs --tail 25 "$CONTAINER_NAME" >&2 || true
    die "node failed to start. Full logs: docker logs $CONTAINER_NAME"
  fi
  providers="$(curl -sf -m 10 "$COORD_URL/v1/providers" 2>/dev/null || true)"
  if [[ -n "$NODE_ID" && "$providers" == *"$NODE_ID"* ]] || [[ -z "$NODE_ID" && "$providers" == *"\"$NODE_NAME\""* ]]; then
    joined=1
    break
  fi
  sleep 5
done

if (( ! joined )); then
  warn "node did not appear at $COORD_URL/v1/providers within 3 minutes."
  warn "it may still be loading the model. Check:"
  warn "  docker logs -f $CONTAINER_NAME"
  warn "  curl -s $COORD_URL/v1/providers"
  if [[ -f "$FEYA_HOME/state/claimed.json" ]]; then
    warn "(pairing itself succeeded — the account binding is saved)"
  fi
  exit 1
fi

cat <<DONE

===============================================================
Your box is on the Feya network.

  Node:       $NODE_NAME (id: ${NODE_ID:-unknown})
  Container:  $CONTAINER_NAME  (restarts automatically on reboot)
  Dashboard:  $COORD_URL
  Logs:       docker logs -f $CONTAINER_NAME
  Verify:     curl -s $COORD_URL/v1/providers | grep -o '"node_name":"[^"]*"'

  Pause:      docker stop $CONTAINER_NAME
  Leave:      docker rm -f $CONTAINER_NAME
              (keep $FEYA_HOME/state to rejoin with the same identity)
===============================================================
DONE
