#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Deploy NVIDIA Isaac Sim (container, headless) to a remote Ubuntu host.

Usage:
  scripts/deploy_isaac_remote.sh --host root@86.38.238.177 [options]

Required:
  --host USER@IP            Remote SSH target.

Options:
  --ssh-key PATH            SSH private key (default: ~/.ssh/id_rsa)
  --image TAG               Isaac Sim image tag (default: 5.1.0)
  --container NAME          Container name (default: isaac-sim)
  --remote-root PATH        Remote cache root (default: /root/docker/isaac-sim)
  --nvcr-key KEY            NGC/NVCR API key (default: NVCR_APIKEY env or ./.env)
  --public-ip IP            Public IP advertised to WebRTC client (default: auto-detect on remote)
  --allow-ip IP             Restrict UFW to this source IP for ports 22/49100 and 47998/udp
  --lockdown                Enable UFW and apply lock rules (requires --allow-ip)
  --skip-compat-check       Skip compatibility checker step
  --help                    Show this help

Examples:
  scripts/deploy_isaac_remote.sh --host root@86.38.238.177
  scripts/deploy_isaac_remote.sh --host root@1.2.3.4 --lockdown --allow-ip 5.29.17.220

Notes:
  - This script installs Docker and NVIDIA Container Toolkit if missing.
  - It logs into nvcr.io, pulls nvcr.io/nvidia/isaac-sim:<tag>,
    runs compatibility check, and starts ./runheadless.sh -v in detached mode.
EOF
}

HOST=""
SSH_KEY="${HOME}/.ssh/id_rsa"
IMAGE_TAG="5.1.0"
CONTAINER_NAME="isaac-sim"
REMOTE_ROOT="/root/docker/isaac-sim"
NVCR_KEY="${NVCR_APIKEY:-}"
PUBLIC_IP=""
ALLOW_IP=""
LOCKDOWN="false"
SKIP_COMPAT="false"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --host)
      HOST="$2"
      shift 2
      ;;
    --ssh-key)
      SSH_KEY="$2"
      shift 2
      ;;
    --image)
      IMAGE_TAG="$2"
      shift 2
      ;;
    --container)
      CONTAINER_NAME="$2"
      shift 2
      ;;
    --remote-root)
      REMOTE_ROOT="$2"
      shift 2
      ;;
    --nvcr-key)
      NVCR_KEY="$2"
      shift 2
      ;;
    --public-ip)
      PUBLIC_IP="$2"
      shift 2
      ;;
    --allow-ip)
      ALLOW_IP="$2"
      shift 2
      ;;
    --lockdown)
      LOCKDOWN="true"
      shift
      ;;
    --skip-compat-check)
      SKIP_COMPAT="true"
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      echo "Unknown arg: $1" >&2
      usage
      exit 1
      ;;
  esac
done

if [[ -z "$HOST" ]]; then
  echo "Error: --host is required" >&2
  usage
  exit 1
fi

if [[ ! -f "$SSH_KEY" ]]; then
  echo "Error: SSH key not found: $SSH_KEY" >&2
  exit 1
fi

if [[ -z "$NVCR_KEY" && -f ./.env ]]; then
  # shellcheck disable=SC1091
  source ./.env
  NVCR_KEY="${NVCR_APIKEY:-}"
fi

if [[ -z "$NVCR_KEY" ]]; then
  echo "Error: missing NVCR API key. Set NVCR_APIKEY, pass --nvcr-key, or provide it in ./.env" >&2
  exit 1
fi

if [[ "$LOCKDOWN" == "true" && -z "$ALLOW_IP" ]]; then
  echo "Error: --lockdown requires --allow-ip" >&2
  exit 1
fi

SSH_OPTS=(
  -i "$SSH_KEY"
  -o BatchMode=yes
  -o ConnectTimeout=20
  -o StrictHostKeyChecking=no
  -o UserKnownHostsFile=/dev/null
)

echo "==> Checking SSH connectivity"
ssh "${SSH_OPTS[@]}" "$HOST" "echo 'Connected to:' \"\$(hostname -f 2>/dev/null || hostname)\""

if [[ -z "$PUBLIC_IP" ]]; then
  PUBLIC_IP="$(ssh "${SSH_OPTS[@]}" "$HOST" "curl -s ifconfig.me || true")"
fi

if [[ -z "$PUBLIC_IP" ]]; then
  echo "Error: could not auto-detect public IP. Pass --public-ip." >&2
  exit 1
fi

echo "==> Deploying Isaac Sim ${IMAGE_TAG} on ${HOST}"
ssh "${SSH_OPTS[@]}" "$HOST" \
  "NVCR_KEY='${NVCR_KEY}' IMAGE_TAG='${IMAGE_TAG}' CONTAINER_NAME='${CONTAINER_NAME}' REMOTE_ROOT='${REMOTE_ROOT}' LOCKDOWN='${LOCKDOWN}' ALLOW_IP='${ALLOW_IP}' SKIP_COMPAT='${SKIP_COMPAT}' PUBLIC_IP='${PUBLIC_IP}' bash -s" <<'REMOTE'
set -euo pipefail

if [[ "$(id -u)" -ne 0 ]]; then
  echo "Run this remote setup as root." >&2
  exit 1
fi

echo "[1/8] OS / GPU summary"
lsb_release -ds || true
uname -r || true
nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader || true

echo "[2/8] Install Docker if missing"
if ! command -v docker >/dev/null 2>&1; then
  curl -fsSL https://get.docker.com -o /tmp/get-docker.sh
  sh /tmp/get-docker.sh
fi

echo "[3/8] Install NVIDIA container toolkit if missing"
if ! dpkg -s nvidia-container-toolkit >/dev/null 2>&1; then
  apt-get update
  apt-get install -y ca-certificates curl gnupg
  install -m 0755 -d /etc/apt/keyrings
  curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /etc/apt/keyrings/nvidia-container-toolkit-keyring.gpg
  chmod a+r /etc/apt/keyrings/nvidia-container-toolkit-keyring.gpg
  curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
    | sed 's#deb https://#deb [signed-by=/etc/apt/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
    > /etc/apt/sources.list.d/nvidia-container-toolkit.list
  apt-get update
  apt-get install -y nvidia-container-toolkit
fi

nvidia-ctk runtime configure --runtime=docker
systemctl restart docker

echo "[4/8] Verify Docker GPU passthrough"
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi >/dev/null

echo "[5/8] Login to nvcr.io and pull image"
printf '%s' "${NVCR_KEY}" | docker login nvcr.io -u '$oauthtoken' --password-stdin
docker pull "nvcr.io/nvidia/isaac-sim:${IMAGE_TAG}"

echo "[6/8] Create persistent cache/config/log directories"
mkdir -p \
  "${REMOTE_ROOT}/cache/main/ov" \
  "${REMOTE_ROOT}/cache/main/warp" \
  "${REMOTE_ROOT}/cache/computecache" \
  "${REMOTE_ROOT}/config" \
  "${REMOTE_ROOT}/data/documents" \
  "${REMOTE_ROOT}/data/Kit" \
  "${REMOTE_ROOT}/logs" \
  "${REMOTE_ROOT}/pkg"
chown -R 1234:1234 "${REMOTE_ROOT}"

if [[ "${SKIP_COMPAT}" != "true" ]]; then
  echo "[7/8] Run compatibility checker"
  docker run --rm --gpus all --network=host --entrypoint bash \
    "nvcr.io/nvidia/isaac-sim:${IMAGE_TAG}" \
    ./isaac-sim.compatibility_check.sh --/app/quitAfter=10 --no-window
fi

echo "[8/8] Start Isaac Sim headless container"
docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true
docker run -d \
  --name "${CONTAINER_NAME}" \
  --entrypoint bash \
  --gpus all \
  --network=host \
  -e ACCEPT_EULA=Y \
  -e PRIVACY_CONSENT=Y \
  -v "${REMOTE_ROOT}/cache/main:/isaac-sim/.cache:rw" \
  -v "${REMOTE_ROOT}/cache/computecache:/isaac-sim/.nv/ComputeCache:rw" \
  -v "${REMOTE_ROOT}/logs:/isaac-sim/.nvidia-omniverse/logs:rw" \
  -v "${REMOTE_ROOT}/config:/isaac-sim/.nvidia-omniverse/config:rw" \
  -v "${REMOTE_ROOT}/data:/isaac-sim/.local/share/ov/data:rw" \
  -v "${REMOTE_ROOT}/pkg:/isaac-sim/.local/share/ov/pkg:rw" \
  -u 1234:1234 \
  "nvcr.io/nvidia/isaac-sim:${IMAGE_TAG}" \
  -lc './runheadless.sh -v --/exts/omni.kit.livestream.app/primaryStream/publicIp='"${PUBLIC_IP}"' --/exts/omni.kit.livestream.app/primaryStream/signalPort=49100 --/exts/omni.kit.livestream.app/primaryStream/streamPort=47998'

echo "Waiting for app readiness (up to 180s)"
for _ in $(seq 1 60); do
  if docker logs "${CONTAINER_NAME}" 2>&1 | grep -q "app ready"; then
    break
  fi
  sleep 3
done

if docker logs "${CONTAINER_NAME}" 2>&1 | grep -q "app ready"; then
  echo "Isaac Sim app is ready."
else
  echo "Warning: app ready marker not seen yet. Check: docker logs -f ${CONTAINER_NAME}" >&2
fi

if [[ "${LOCKDOWN}" == "true" ]]; then
  echo "Applying UFW lockdown for ${ALLOW_IP}"
  apt-get update
  apt-get install -y ufw
  ufw --force reset
  ufw default deny incoming
  ufw default allow outgoing
  ufw allow from "${ALLOW_IP}" to any port 22 proto tcp
  ufw allow from "${ALLOW_IP}" to any port 49100 proto tcp
  ufw allow from "${ALLOW_IP}" to any port 47998 proto udp
  ufw allow from "${ALLOW_IP}" to any port 8011 proto tcp
  ufw --force enable
fi

echo "Open/listening ports snapshot:"
ss -tulpen | grep -E ':49100|:47998|:8011' || true
echo "UFW status:"
ufw status verbose || true
REMOTE

echo "==> Done"
echo "Connect target: ${HOST#*@}"
echo "- Streaming port: 49100"
echo "- Media port:      47998/udp"
echo "- Service/docs:    http://${HOST#*@}:8011/docs"
echo "Useful checks:"
echo "  ssh -i ${SSH_KEY} ${HOST} 'docker ps --filter name=${CONTAINER_NAME}'"
echo "  ssh -i ${SSH_KEY} ${HOST} 'docker logs -f ${CONTAINER_NAME}'"