#!/bin/bash
# provision-sd.sh — flash + prepare a JetPack 6.2 SD card for a Dragon NVR unit
# (Jetson Orin Nano devkit), with zero first-boot interaction.
#
# What it does, in order:
#   1. Downloads the JetPack 6.2 SD-card image zip (resumable, cached) if missing.
#   2. Writes it to the target SD card (streamed straight out of the zip).
#   3. Expands the APP (rootfs) partition + filesystem to fill the card —
#      the resize step oem-config would normally offer (skip with --no-expand).
#   4. Mounts the APP partition and injects, offline:
#        - a login user (passwd/shadow/group edits — no chroot, no qemu needed)
#        - root password and/or SSH authorized_keys; PermitRootLogin
#        - oem-config removal (the Ubuntu first-boot wizard never runs)
#        - default boot target = multi-user (text console, no GUI)
#        - hostname
#        - a first-boot systemd oneshot that, on the Jetson itself:
#            * switches nvpmodel to the MAXN/MAXN_SUPER power mode
#            * formats a blank NVMe SSD (GPT + ext4, label "vs1") and mounts
#              it at /mnt/vs1 via fstab (existing vs1 disks are reused,
#              disks with foreign data are left untouched)
#            * generates root's ssh keypair (/root/.ssh/id_rsa) — the
#              identity dragon's reverse-ssh tunnel authenticates with;
#              generated on-device so the private key never leaves the unit
#            * reboots once so the power mode takes effect
#        - a docker-setup systemd oneshot that runs on the boot after that,
#          waits for network, then installs docker.io + compose and configures
#          the NVIDIA container runtime as docker's default; retries on every
#          boot until it succeeds (skip with --skip-docker)
#        - a per-boot jetson_clocks service (clocks pinned to max)
#
# Dragon itself is NOT installed — this only prepares the OS.
#
# Typical use (passwords and ssh key have shop defaults — see Options):
#   sudo ./provision-sd.sh --device /dev/sdX --hostname cloudnvr
#
# Requires (on the machine writing the card): bash, curl, unzip, openssl,
# blkid, partprobe (parted), lsblk. Must run as root.

set -euo pipefail

# ---------------------------------------------------------------- defaults --
IMAGE_URL="https://developer.download.nvidia.com/embedded/L4T/r36_Release_v4.3/jp62-orin-nano-sd-card-image.zip"
CACHE_DIR="${HOME}/.cache/dragon-provision"
DEVICE=""
LOGIN_USER="dragon"
LOGIN_PASS="drb2245"           # standard shop password; override with --pass
ROOT_PASS="drb2245"            # override with --root-pass
SSH_KEY_FILE=""                # defaults to ~/.ssh/id_ed25519.pub when present (see below)
HOSTNAME_NEW="cloudnvr"
ASSUME_YES=0
NO_FLASH=0
DOWNLOAD_ONLY=0
SKIP_DOCKER=0
NO_EXPAND=0

usage() {
    sed -n '2,20p' "$0" | sed 's/^# \{0,1\}//'
    cat <<EOF

Options:
  --device /dev/sdX     Target SD card block device (whole disk, not a partition)
  --user NAME           Login user to create (default: ${LOGIN_USER})
  --pass PASS           Password for the login user (default: ${LOGIN_PASS})
  --root-pass PASS      Root password; also enables PermitRootLogin yes (default: ${ROOT_PASS})
  --ssh-key FILE        Public key installed for root and the login user
                        (default: ~/.ssh/id_ed25519.pub when present)
  --hostname NAME       Hostname (default: ${HOSTNAME_NEW})
  --image-url URL       SD image zip (default: JetPack 6.2 / L4T 36.4.3)
  --cache-dir DIR       Download cache (default: ${CACHE_DIR})
  --skip-docker         Don't install the docker + NVIDIA runtime first-boot unit
  --no-expand           Don't grow the APP partition/filesystem to fill the card
  --no-flash            Skip download+write; only inject into an already-flashed card
  --download-only       Fetch/verify the image zip and exit
  --yes                 Skip the "type FLASH" confirmation
  -h, --help            This help
EOF
    exit "${1:-0}"
}

die()  { echo "ERROR: $*" >&2; exit 1; }
info() { echo "==> $*"; }

# ------------------------------------------------------------------- args --
while [ $# -gt 0 ]; do
    case "$1" in
        --device)        DEVICE="$2"; shift 2 ;;
        --user)          LOGIN_USER="$2"; shift 2 ;;
        --pass)          LOGIN_PASS="$2"; shift 2 ;;
        --root-pass)     ROOT_PASS="$2"; shift 2 ;;
        --ssh-key)       SSH_KEY_FILE="$2"; shift 2 ;;
        --hostname)      HOSTNAME_NEW="$2"; shift 2 ;;
        --image-url)     IMAGE_URL="$2"; shift 2 ;;
        --cache-dir)     CACHE_DIR="$2"; shift 2 ;;
        --skip-docker)   SKIP_DOCKER=1; shift ;;
        --no-expand)     NO_EXPAND=1; shift ;;
        --no-flash)      NO_FLASH=1; shift ;;
        --download-only) DOWNLOAD_ONLY=1; shift ;;
        --yes)           ASSUME_YES=1; shift ;;
        -h|--help)       usage 0 ;;
        *)               echo "Unknown option: $1" >&2; usage 1 ;;
    esac
done

[ "$(id -u)" -eq 0 ] || die "must run as root (dd/mount)"
for tool in curl unzip openssl blkid partprobe lsblk parted e2fsck resize2fs; do
    command -v "$tool" >/dev/null || die "missing tool: $tool"
done
echo "$LOGIN_USER" | grep -qE '^[a-z_][a-z0-9_-]*$' || die "invalid --user '$LOGIN_USER'"

# Soft ssh-key default: install the invoking user's ed25519 pubkey when it
# exists; skip quietly when it doesn't (an explicit --ssh-key still dies on
# an unreadable file below).
if [ -z "$SSH_KEY_FILE" ]; then
    if [ -r "${HOME}/.ssh/id_ed25519.pub" ]; then
        SSH_KEY_FILE="${HOME}/.ssh/id_ed25519.pub"
        info "using default ssh key ${SSH_KEY_FILE}"
    else
        info "no --ssh-key and ${HOME}/.ssh/id_ed25519.pub not found - skipping key install"
    fi
fi
if [ -n "$SSH_KEY_FILE" ]; then
    [ -r "$SSH_KEY_FILE" ] || die "cannot read --ssh-key $SSH_KEY_FILE"
    grep -qE '^(ssh|ecdsa)-' "$SSH_KEY_FILE" || die "$SSH_KEY_FILE does not look like a public key"
fi

# Require --device up front (unless only downloading) so a mistyped
# invocation fails immediately instead of kicking off an ~11 GB download.
if [ "$DOWNLOAD_ONLY" -eq 0 ] && [ -z "$DEVICE" ]; then
    die "--device is required"
fi

# --------------------------------------------------------------- download --
ZIP_FILE="${CACHE_DIR}/$(basename "$IMAGE_URL")"

download_image() {
    mkdir -p "$CACHE_DIR"
    local remote_len
    remote_len=$(curl -sIL --max-time 60 "$IMAGE_URL" | tr -d '\r' \
                 | awk 'tolower($1)=="content-length:" {len=$2} END {print len}')
    if [ -s "$ZIP_FILE" ] && [ -n "$remote_len" ] \
       && [ "$(stat -c %s "$ZIP_FILE")" = "$remote_len" ]; then
        info "image already cached: $ZIP_FILE ($remote_len bytes)"
        return
    fi
    info "downloading $IMAGE_URL (~11 GB, resumable) ..."
    curl -L -C - --fail -o "$ZIP_FILE" "$IMAGE_URL"
    if [ -n "$remote_len" ] && [ "$(stat -c %s "$ZIP_FILE")" != "$remote_len" ]; then
        die "download size mismatch (got $(stat -c %s "$ZIP_FILE"), expected $remote_len)"
    fi
    info "download complete"
}

if [ "$NO_FLASH" -eq 0 ]; then
    download_image
fi
[ "$DOWNLOAD_ONLY" -eq 1 ] && { info "done (--download-only)"; exit 0; }

# ----------------------------------------------------------- device checks --
[ -n "$DEVICE" ] || die "--device is required"
[ -b "$DEVICE" ] || die "$DEVICE is not a block device"
# refuse partitions ("/dev/sdb1") — we need the whole disk (loop allowed for testing)
case "$(lsblk -ndo TYPE "$DEVICE")" in
    disk|loop) ;;
    *) die "$DEVICE is not a whole disk" ;;
esac

ROOT_SRC=$(findmnt -no SOURCE / || true)
case "$ROOT_SRC" in
    "$DEVICE"*) die "$DEVICE holds the running system — refusing" ;;
esac
if lsblk -no MOUNTPOINT "$DEVICE" | grep -q .; then
    die "$DEVICE has mounted partitions — unmount first"
fi

# partition name prefix: /dev/sdb -> sdb1, /dev/mmcblk0 -> mmcblk0p1
PART_SEP=""
echo "$DEVICE" | grep -qE '[0-9]$' && PART_SEP="p"
APP_PART="${DEVICE}${PART_SEP}1"

if [ "$ASSUME_YES" -eq 0 ]; then
    echo
    lsblk -o NAME,SIZE,MODEL,TRAN,MOUNTPOINT "$DEVICE"
    echo
    if [ "$NO_FLASH" -eq 0 ]; then
        echo "ALL DATA ON $DEVICE WILL BE DESTROYED."
        printf 'Type FLASH to continue: '
        read -r answer
        [ "$answer" = "FLASH" ] || die "aborted"
    else
        printf 'Inject configuration into %s? [y/N] ' "$APP_PART"
        read -r answer
        [ "$answer" = "y" ] || [ "$answer" = "Y" ] || die "aborted"
    fi
fi

# ---------------------------------------------------------------- passwords --
if [ -z "$LOGIN_PASS" ]; then
    [ -t 0 ] || die "--pass required when not running interactively"
    printf 'Password for user %s: ' "$LOGIN_USER"
    read -rs LOGIN_PASS; echo
    [ -n "$LOGIN_PASS" ] || die "empty password"
fi
if [ -z "$ROOT_PASS" ] && [ -z "$SSH_KEY_FILE" ]; then
    echo "WARNING: no --root-pass and no --ssh-key — root will not be reachable over SSH." >&2
fi

hash_pw() { openssl passwd -6 "$1"; }
LOGIN_HASH=$(hash_pw "$LOGIN_PASS")
ROOT_HASH=""
[ -n "$ROOT_PASS" ] && ROOT_HASH=$(hash_pw "$ROOT_PASS")
DAYS_SINCE_EPOCH=$(( $(date +%s) / 86400 ))

# -------------------------------------------------------------------- flash --
if [ "$NO_FLASH" -eq 0 ]; then
    IMG_NAME=$(unzip -Z1 "$ZIP_FILE" | grep -m1 '\.img$') \
        || die "no .img inside $ZIP_FILE"
    info "writing $IMG_NAME -> $DEVICE (streamed from zip; takes a while) ..."
    unzip -p "$ZIP_FILE" "$IMG_NAME" \
        | dd of="$DEVICE" bs=16M conv=fsync status=progress
    sync
fi

info "re-reading partition table"
partprobe "$DEVICE"
for _ in $(seq 1 20); do
    [ -b "$APP_PART" ] && break
    sleep 0.5
done
[ -b "$APP_PART" ] || die "APP partition $APP_PART did not appear"
FSTYPE=$(blkid -o value -s TYPE "$APP_PART" || true)
[ "$FSTYPE" = "ext4" ] || die "$APP_PART is not ext4 (got '$FSTYPE') — wrong image/device?"

# ------------------------------------------------------------------ expand --
# The image's GPT covers only the image size, so first let parted relocate the
# backup GPT header to the true end of the card (-f answers the fix prompt),
# then grow APP. The Jetson SD layout puts APP (partition 1) physically last
# on disk precisely so this in-place grow works — the same resize oem-config
# offers on first boot. Idempotent: a second run is a no-op.
if [ "$NO_EXPAND" -eq 0 ]; then
    info "expanding APP partition + filesystem to fill the card"
    parted -sf "$DEVICE" print >/dev/null 2>&1 || true
    parted -s "$DEVICE" resizepart 1 100% \
        || die "resizepart failed (is APP not the last partition on disk?) — rerun with --no-expand to skip"
    partprobe "$DEVICE"; sleep 1
    set +e
    e2fsck -pf "$APP_PART"; fsck_rc=$?
    set -e
    [ "$fsck_rc" -le 2 ] || die "e2fsck failed on $APP_PART (rc=$fsck_rc)"
    resize2fs "$APP_PART"
fi

# ------------------------------------------------------------------ inject --
MNT=$(mktemp -d /tmp/dragon-sd.XXXXXX)
cleanup() { umount "$MNT" 2>/dev/null || true; rmdir "$MNT" 2>/dev/null || true; }
trap cleanup EXIT
mount "$APP_PART" "$MNT"
[ -f "$MNT/etc/passwd" ] || die "$APP_PART does not look like a rootfs"

R="$MNT"

info "creating user '$LOGIN_USER' (uid 1000)"
if grep -q "^${LOGIN_USER}:" "$R/etc/passwd"; then
    info "  user already exists in image — skipping creation"
else
    grep -q ':1000:' "$R/etc/passwd" && die "uid 1000 already taken in image"
    echo "${LOGIN_USER}:x:1000:1000:${LOGIN_USER}:/home/${LOGIN_USER}:/bin/bash" >> "$R/etc/passwd"
    echo "${LOGIN_USER}:${LOGIN_HASH}:${DAYS_SINCE_EPOCH}:0:99999:7:::"          >> "$R/etc/shadow"
    echo "${LOGIN_USER}:x:1000:"    >> "$R/etc/group"
    echo "${LOGIN_USER}:!::"        >> "$R/etc/gshadow"
    cp -a "$R/etc/skel" "$R/home/${LOGIN_USER}"
    chown -R 1000:1000 "$R/home/${LOGIN_USER}"
fi

# supplementary groups (only those that exist in the image)
for grp in sudo adm video render audio plugdev dialout i2c gpio; do
    grep -q "^${grp}:" "$R/etc/group" || continue
    grep -qE "^${grp}:.*[:,]${LOGIN_USER}(,|$)" "$R/etc/group" && continue
    # empty member list ends in ':' -> append name; otherwise append ',name'
    sed -i -E "/^${grp}:/{s/:\$/:${LOGIN_USER}/; t; s/\$/,${LOGIN_USER}/}" "$R/etc/group"
done

if [ -n "$ROOT_HASH" ]; then
    info "setting root password"
    # replace root's password field regardless of what it currently is
    sed -i -E "s|^root:[^:]*:|root:${ROOT_HASH//&/\\&}:|" "$R/etc/shadow"
fi

info "configuring SSH"
mkdir -p "$R/etc/ssh/sshd_config.d"
{
    if [ -n "$ROOT_HASH" ]; then
        echo "PermitRootLogin yes"
    else
        echo "PermitRootLogin prohibit-password"
    fi
} > "$R/etc/ssh/sshd_config.d/60-dragon.conf"

if [ -n "$SSH_KEY_FILE" ]; then
    for home in "$R/root" "$R/home/${LOGIN_USER}"; do
        mkdir -p "$home/.ssh"
        cat "$SSH_KEY_FILE" >> "$home/.ssh/authorized_keys"
        chmod 700 "$home/.ssh"
        chmod 600 "$home/.ssh/authorized_keys"
    done
    chown -R 1000:1000 "$R/home/${LOGIN_USER}/.ssh"
fi
# make sure sshd is enabled
mkdir -p "$R/etc/systemd/system/multi-user.target.wants"
[ -f "$R/lib/systemd/system/ssh.service" ] && \
    ln -sf /lib/systemd/system/ssh.service \
        "$R/etc/systemd/system/multi-user.target.wants/ssh.service"

info "disabling first-boot wizard (oem-config)"
find "$R/etc/systemd" \( -name '*oem-config*' -o -name '*nv-oem-config*' \) \
     -print -exec rm -rf {} + 2>/dev/null || true
rm -f "$R/var/lib/oem-config/run"

info "boot to text console (multi-user.target)"
ln -sf /lib/systemd/system/multi-user.target "$R/etc/systemd/system/default.target"

info "setting hostname '$HOSTNAME_NEW'"
echo "$HOSTNAME_NEW" > "$R/etc/hostname"
if grep -q '^127\.0\.1\.1' "$R/etc/hosts"; then
    sed -i "s/^127\.0\.1\.1.*/127.0.1.1\t${HOSTNAME_NEW}/" "$R/etc/hosts"
else
    printf '127.0.1.1\t%s\n' "$HOSTNAME_NEW" >> "$R/etc/hosts"
fi

# --------------------------------------------------- first-boot provisioning --
info "installing first-boot service (nvpmodel MAXN, NVMe -> /mnt/vs1, root ssh keypair)"

install -d -m 755 "$R/usr/local/sbin"
cat > "$R/usr/local/sbin/dragon-firstboot.sh" <<'FIRSTBOOT'
#!/bin/bash
# Dragon NVR first-boot provisioning. Runs once on the Jetson, then reboots.
# Idempotent: guarded by /var/lib/dragon-firstboot.done.
set -u
exec >>/var/log/dragon-firstboot.log 2>&1
echo "=== dragon-firstboot $(date -Is) ==="
set -x

MARKER=/var/lib/dragon-firstboot.done
[ -e "$MARKER" ] && exit 0

# --- 1. max power mode -------------------------------------------------------
# Find the MAXN (or MAXN_SUPER) mode id in nvpmodel.conf rather than hardcoding
# it — the id differs between JetPack releases and modules.
if [ -r /etc/nvpmodel.conf ] && command -v nvpmodel >/dev/null; then
    maxn_id=$(grep -oE 'ID=[0-9]+[[:space:]]+NAME=[A-Z0-9_]*MAXN[A-Z0-9_]*' /etc/nvpmodel.conf \
              | grep -oE 'ID=[0-9]+' | cut -d= -f2 | tail -1)
    if [ -n "${maxn_id:-}" ]; then
        # nvpmodel asks "reboot now? [Y/n]" for modes that need one; answer no —
        # we reboot ourselves at the end.
        printf 'n\n' | nvpmodel -m "$maxn_id" || true
    fi
fi

# --- 2. NVMe SSD -> /mnt/vs1 -------------------------------------------------
mkdir -p /mnt/vs1
part=$(blkid -L vs1 2>/dev/null || true)
if [ -z "$part" ]; then
    nvme=""
    for d in /dev/nvme0n1 /dev/nvme1n1; do
        [ -b "$d" ] && nvme="$d" && break
    done
    if [ -n "$nvme" ]; then
        nparts=$(lsblk -no NAME "$nvme" | wc -l)     # 1 == disk only, no partitions
        sig=$(blkid -o value -s TYPE "$nvme" 2>/dev/null || true)
        if [ "$nparts" -eq 1 ] && [ -z "$sig" ]; then
            parted -s "$nvme" mklabel gpt mkpart vs1 ext4 1MiB 100%
            udevadm settle
            partprobe "$nvme"; sleep 2
            mkfs.ext4 -q -F -L vs1 "${nvme}p1"
            part="${nvme}p1"
        else
            echo "NVMe $nvme has existing partitions/filesystem and no 'vs1' label — leaving it alone"
        fi
    else
        echo "no NVMe SSD found — /mnt/vs1 will stay on the SD card"
    fi
fi
if [ -n "$part" ]; then
    if ! grep -q '^LABEL=vs1[[:space:]]' /etc/fstab; then
        echo 'LABEL=vs1 /mnt/vs1 ext4 defaults,noatime,nofail,x-systemd.device-timeout=10 0 2' >> /etc/fstab
    fi
    systemctl daemon-reload
    mountpoint -q /mnt/vs1 || mount /mnt/vs1 || true
fi

# --- 3. device ssh keypair (reverse-ssh tunnel identity) ----------------------
# dragon's PMSG_REQUEST_REVERSE_SSH handler opens the tunnel as root with the
# default identity and uploads /root/.ssh/id_rsa.pub to the backend
# (PMSG_SSH_PUB_KEY) when the tunnel host doesn't know us yet — so the key
# must be RSA at that exact path. Generated here, on the device, so the
# private key never exists off-unit and cloned SD images can't share one.
if [ ! -f /root/.ssh/id_rsa ]; then
    mkdir -p -m 700 /root/.ssh
    ssh-keygen -t rsa -b 4096 -N "" -C "cloudnvr-$(hostname)" -f /root/.ssh/id_rsa
fi

# --- 4. done — reboot once so the nvpmodel mode is live ----------------------
touch "$MARKER"
systemctl disable dragon-firstboot.service || true
echo "=== dragon-firstboot complete, rebooting ==="
systemctl reboot
FIRSTBOOT
chmod 755 "$R/usr/local/sbin/dragon-firstboot.sh"

cat > "$R/etc/systemd/system/dragon-firstboot.service" <<'UNIT'
[Unit]
Description=Dragon NVR first-boot provisioning (power mode, NVMe /mnt/vs1, ssh keypair)
ConditionPathExists=!/var/lib/dragon-firstboot.done
After=local-fs.target

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/dragon-firstboot.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
UNIT
ln -sf /etc/systemd/system/dragon-firstboot.service \
    "$R/etc/systemd/system/multi-user.target.wants/dragon-firstboot.service"

# jetson_clocks every boot (pins clocks at the max for the active nvpmodel mode)
cat > "$R/etc/systemd/system/dragon-clocks.service" <<'UNIT'
[Unit]
Description=Pin Jetson clocks to maximum (jetson_clocks)
# only after first-boot provisioning has set the final nvpmodel mode
ConditionPathExists=/usr/bin/jetson_clocks
ConditionPathExists=/var/lib/dragon-firstboot.done
After=nvpmodel.service

[Service]
Type=oneshot
ExecStart=/usr/bin/jetson_clocks
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
UNIT
ln -sf /etc/systemd/system/dragon-clocks.service \
    "$R/etc/systemd/system/multi-user.target.wants/dragon-clocks.service"

# Periodic system-clock save. The Orin Nano has no RTC battery, so after a
# power cut the clock comes up at the 1970 epoch and systemd-timesyncd restores
# it from the mtime of /var/lib/systemd/timesync/clock — which JetPack's
# systemd 249 only touches on a successful NTP sync (poll interval backs off to
# 34 min) or a clean shutdown. Touching it every minute caps the post-power-cut
# clock regression at ~60 s even with NTP unreachable, so new recordings can't
# collide with pre-cut ones. (Backport of systemd 250's SaveIntervalSec=60
# default; ~16 KB/min of ext4 metadata writes, negligible SD wear.)
cat > "$R/etc/systemd/system/dragon-clocksave.service" <<'UNIT'
[Unit]
Description=Save system clock for boot-time restore (no RTC battery)
ConditionPathExists=/var/lib/systemd/timesync/clock

[Service]
Type=oneshot
ExecStart=/usr/bin/touch -c /var/lib/systemd/timesync/clock
UNIT

cat > "$R/etc/systemd/system/dragon-clocksave.timer" <<'UNIT'
[Unit]
Description=Save system clock every minute (no RTC battery)

[Timer]
OnBootSec=60
OnUnitActiveSec=60

[Install]
WantedBy=timers.target
UNIT
mkdir -p "$R/etc/systemd/system/timers.target.wants"
ln -sf /etc/systemd/system/dragon-clocksave.timer \
    "$R/etc/systemd/system/timers.target.wants/dragon-clocksave.timer"

# ------------------------------------------------- docker + nvidia runtime --
if [ "$SKIP_DOCKER" -eq 0 ]; then
    info "installing docker-setup service (docker.io + compose + nvidia runtime)"

    cat > "$R/usr/local/sbin/dragon-docker-setup.sh" <<'DOCKERSETUP'
#!/bin/bash
# Dragon NVR docker setup. Runs on the boot after dragon-firstboot's reboot;
# needs network, so on failure it exits without the marker and the unit
# retries on every subsequent boot until it succeeds.
set -u
LOG_FILE=${DRAGON_DOCKER_SETUP_LOG:-/var/log/dragon-docker-setup.log}
exec >>"$LOG_FILE" 2>&1
echo "=== dragon-docker-setup $(date -Is) ==="
set -x

MARKER=${DRAGON_DOCKER_SETUP_MARKER:-/var/lib/dragon-docker-setup.done}
MIN_VALID_EPOCH=1704067200
[ -e "$MARKER" ] && exit 0

clock_is_synchronized() {
    now=$(date +%s)
    [ "$now" -ge "$MIN_VALID_EPOCH" ] \
        && [ "$(timedatectl show -p NTPSynchronized --value 2>/dev/null)" = yes ]
}

# Bounded wait for a network and a clock apt can use (up to ~5 min this boot).
ok=0
for _ in $(seq 1 60); do
    if clock_is_synchronized \
        && curl -fsI --max-time 5 http://ports.ubuntu.com >/dev/null 2>&1; then
        ok=1; break
    fi
    sleep 5
done
if [ "$ok" != 1 ]; then
    echo "no network or valid clock - will retry next boot"
    exit 1
fi

export DEBIAN_FRONTEND=noninteractive
apt-get update || exit 1
if ! command -v docker >/dev/null 2>&1; then
    apt-get install -y docker.io || exit 1
fi
# compose v2 plugin ("docker compose"): in jammy universe on current point
# releases.
if ! docker compose version >/dev/null 2>&1; then
    apt-get install -y docker-compose-v2 || exit 1
fi

# NVIDIA container toolkit: recent JetPack 6 images ship it; otherwise it
# comes from the L4T apt repo (package name varies across releases).
if ! command -v nvidia-ctk >/dev/null; then
    apt-get install -y nvidia-container-toolkit \
        || apt-get install -y nvidia-container \
        || exit 1
fi
if command -v nvidia-ctk >/dev/null; then
    # register the nvidia runtime with docker and make it the default
    # (jetson convention); older toolkits lack --set-as-default
    nvidia-ctk runtime configure --runtime=docker --set-as-default \
        || nvidia-ctk runtime configure --runtime=docker \
        || exit 1
else
    echo "nvidia-ctk not found - will retry next boot"
    exit 1
fi

systemctl enable docker || exit 1
systemctl restart docker || exit 1

# Do not claim success until the installer prerequisites are usable.
docker compose version >/dev/null 2>&1 || exit 1
docker info >/dev/null 2>&1 || exit 1
if [ "$(docker info --format '{{.DefaultRuntime}}')" != nvidia ]; then
    echo "Docker default runtime is not nvidia - will retry next boot"
    exit 1
fi

# Support tooling must not prevent recorder provisioning after Docker is ready.
# Bound downloads, not dpkg: killing a package transaction can require repair.
if ! tmux -V >/dev/null 2>&1; then
    if ! apt-get install -y --no-install-recommends tmux -o Acquire::http::Timeout=30 -o Acquire::https::Timeout=30 -o Acquire::Retries=2 \
            || ! tmux -V >/dev/null 2>&1; then
        echo "WARNING: host tmux unavailable; nvr update will retry after installation"
    fi
fi

# let the login user run docker without sudo (takes effect on next login)
id -u "@LOGIN_USER@" >/dev/null 2>&1 && usermod -aG docker "@LOGIN_USER@"

touch "$MARKER"
echo "=== dragon-docker-setup complete ==="
exit 0
DOCKERSETUP
    sed -i "s/@LOGIN_USER@/${LOGIN_USER}/g" "$R/usr/local/sbin/dragon-docker-setup.sh"
    chmod 755 "$R/usr/local/sbin/dragon-docker-setup.sh"

    cat > "$R/etc/systemd/system/dragon-docker-setup.service" <<'UNIT'
[Unit]
Description=Dragon NVR docker setup (docker.io + NVIDIA container runtime)
# Runs after dragon-firstboot has completed (i.e. from the second boot on)
# and keeps retrying each boot until its own marker exists.
ConditionPathExists=/var/lib/dragon-firstboot.done
ConditionPathExists=!/var/lib/dragon-docker-setup.done
Wants=network-online.target time-sync.target
After=network-online.target time-sync.target dragon-firstboot.service

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/dragon-docker-setup.sh
RemainAfterExit=yes
TimeoutStartSec=30min

[Install]
WantedBy=multi-user.target
UNIT
    ln -sf /etc/systemd/system/dragon-docker-setup.service \
        "$R/etc/systemd/system/multi-user.target.wants/dragon-docker-setup.service"
fi

sync
info "unmounting"
umount "$MNT"
trap - EXIT
rmdir "$MNT"

cat <<EOF

Done. The card in $DEVICE is ready.

First boot on the Jetson will:
  - come up headless on multi-user.target (no setup wizard, no GUI)
  - run dragon-firstboot: set MAXN power mode, format/mount the NVMe at
    /mnt/vs1 (label 'vs1'), then reboot once
  - after that reboot: sshd up, ${LOGIN_USER}$( [ -n "$ROOT_HASH" ] && echo " and root" ) can log in$( [ "$SKIP_DOCKER" -eq 0 ] && printf '\n  - run dragon-docker-setup once network is up: docker.io + compose +\n    NVIDIA container runtime (default runtime nvidia); retries each boot\n    until it succeeds' )

Logs on the device: /var/log/dragon-firstboot.log$( [ "$SKIP_DOCKER" -eq 0 ] && printf ', /var/log/dragon-docker-setup.log' )
EOF
