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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
IMG="$HOME/Documents/toes-image-builder/toes-ubuntu-base.img"
USBBOOT="$HOME/apps/forks/usbboot"
PROVISION_SCRIPT="${PROVISION_SCRIPT:-$SCRIPT_DIR/scripts/provision-device.sh}"
PROVISION_DEVICE_ID="${PROVISION_DEVICE_ID:-${DEVICE_ID:-device-001}}"
DEFAULT_SERIAL_LOGIN="toes"
DEFAULT_SERIAL_PASSWORD="nicetrywiseguy"
REBOOT_SETTLE_SECS="${REBOOT_SETTLE_SECS:-45}"
SERIAL_BOOT_TIMEOUT_SECS="${SERIAL_BOOT_TIMEOUT_SECS:-180}"

list_external_disks() {
  diskutil list external physical 2>/dev/null | awk '/^\/dev\/disk[0-9]+ / { sub("^/dev/", "", $1); print $1 }' | sort || true
}

disk_bytes() {
  diskutil info "/dev/$1" | awk -F'[()]' '/Disk Size:/ { sub(/ Bytes.*/, "", $2); print $2; exit }'
}

detect_serial_port() {
  local pattern port
  local -a matches

  matches=()
  for pattern in /dev/cu.usbserial\* /dev/cu.usbmodem\*; do
    for port in $pattern; do
      [[ -e "$port" ]] || continue
      matches+=("$port")
    done
  done

  if (( ${#matches[@]} == 1 )); then
    printf '%s\n' "${matches[0]}"
    return 0
  fi

  if (( ${#matches[@]} > 1 )); then
    echo "multiple serial ports detected; set SERIAL_PORT to one of:" >&2
    printf '  %s\n' "${matches[@]}" >&2
    return 2
  fi

  matches=()
  for pattern in /dev/tty.usbserial\* /dev/tty.usbmodem\*; do
    for port in $pattern; do
      [[ -e "$port" ]] || continue
      matches+=("$port")
    done
  done

  if (( ${#matches[@]} == 1 )); then
    printf '%s\n' "${matches[0]}"
    return 0
  fi

  if (( ${#matches[@]} > 1 )); then
    echo "multiple serial ports detected; set SERIAL_PORT to one of:" >&2
    printf '  %s\n' "${matches[@]}" >&2
    return 2
  fi

  return 1
}

resolve_serial_port() {
  local port status

  if [[ -n "${SERIAL_PORT:-}" ]]; then
    [[ -e "$SERIAL_PORT" ]] || { echo "serial port does not exist: $SERIAL_PORT" >&2; return 1; }
    printf '%s\n' "$SERIAL_PORT"
    return 0
  fi

  if port="$(detect_serial_port)"; then
    printf '%s\n' "$port"
    return 0
  fi

  status=$?
  if (( status == 1 )); then
    echo "no serial port detected; set SERIAL_PORT=/dev/cu.*" >&2
  fi
  return 1
}

wait_for_serial_port() {
  local timeout="$1"
  local deadline=$((SECONDS + timeout))
  local port status

  if [[ -n "${SERIAL_PORT:-}" ]]; then
    while (( SECONDS < deadline )); do
      [[ -e "$SERIAL_PORT" ]] && { printf '%s\n' "$SERIAL_PORT"; return 0; }
      sleep 1
    done
    echo "timed out waiting for serial port: $SERIAL_PORT" >&2
    return 1
  fi

  while (( SECONDS < deadline )); do
    if port="$(detect_serial_port 2>/dev/null)"; then
      printf '%s\n' "$port"
      return 0
    fi

    status=$?
    if (( status == 2 )); then
      detect_serial_port >/dev/null || true
      return 1
    fi

    sleep 1
  done

  echo "timed out waiting for serial port; set SERIAL_PORT=/dev/cu.*" >&2
  return 1
}

stop_rpiboot_mass_storage() {
  pkill -f 'rpiboot.*mass-storage-gadget' 2>/dev/null || true
}

reboot_target_over_serial() {
  local port="$1"

  echo "Requesting target reboot over $port"
  stty -f "$port" 115200 raw -echo clocal 2>/dev/null || true
  {
    sleep 0.2
    printf '\r'
    sleep 0.5
    printf 'root\r'
    sleep 0.5
    printf '\r'
    sleep 0.5
    printf 'sync\r'
    sleep 0.5
    printf 'reboot -f\r'
  } >"$port"
}

provision_device() {
  local port="$1"
  local provision_login="${SERIAL_LOGIN:-$DEFAULT_SERIAL_LOGIN}"
  local provision_password="${SERIAL_PASSWORD:-$DEFAULT_SERIAL_PASSWORD}"
  local provision_prompt="${SERIAL_PROMPT:-${provision_login}@}"

  [[ -x "$PROVISION_SCRIPT" ]] || { echo "provision script is not executable: $PROVISION_SCRIPT" >&2; return 1; }

  echo "Provisioning $PROVISION_DEVICE_ID over $port"
  SERIAL_PORT="$port" \
  SERIAL_LOGIN="$provision_login" \
  SERIAL_PASSWORD="$provision_password" \
  SERIAL_PROMPT="$provision_prompt" \
  SERIAL_LOGIN_TIMEOUT_SECS="${SERIAL_LOGIN_TIMEOUT_SECS:-$SERIAL_BOOT_TIMEOUT_SECS}" \
    "$PROVISION_SCRIPT" "$PROVISION_DEVICE_ID"
}

if [[ "${SKIP_BUILD:-0}" != 1 ]]; then
  ssh toes-image-builder 'cd toes-image-builder && ./build.sh && cp build/toes-ubuntu-base.img /mnt/utm/toes-image-builder/'
fi

[[ -f "$IMG" ]] || { echo "missing $IMG" >&2; exit 1; }

target="${DISK:-}"
if [[ -n "$target" ]]; then
  target="${target#/dev/}"
  target="${target#r}"
else
  before="$(mktemp)"
  after="$(mktemp)"
  trap 'rm -f "$before" "$after"' EXIT
  list_external_disks >"$before"
  sudo "$USBBOOT/rpiboot" -d "$USBBOOT/mass-storage-gadget64"

  for _ in {1..30}; do
    list_external_disks >"$after"
    new="$(comm -13 "$before" "$after")"
    [[ -n "$new" ]] && { sleep 2; list_external_disks >"$after"; new="$(comm -13 "$before" "$after")"; break; }
    sleep 1
  done
  [[ -n "${new:-}" ]] || { echo "no new disks" >&2; exit 1; }

  echo "New disks:"
  target_size=0
  for disk in $new; do
    bytes="$(disk_bytes "$disk")"
    echo "  /dev/$disk ${bytes:-?} bytes"
    [[ "$bytes" =~ ^[0-9]+$ ]] && (( bytes > target_size )) && { target="$disk"; target_size="$bytes"; }
  done
fi

[[ "${target:-}" =~ ^disk[0-9]+$ ]] || { echo "bad disk: ${target:-}" >&2; exit 1; }

target_size="$(disk_bytes "$target")"
image_size="$(stat -f %z "$IMG")"
[[ "$target_size" =~ ^[0-9]+$ ]] || { echo "could not size /dev/$target" >&2; exit 1; }
(( target_size >= image_size )) || { echo "/dev/$target is smaller than $IMG" >&2; exit 1; }

echo "Selected /dev/$target"
echo "Image: $IMG"
diskutil info "/dev/$target"
printf 'Type FLASH to erase /dev/%s: ' "$target"
read -r answer
[[ "$answer" == FLASH ]] || { echo "aborted"; exit 1; }

diskutil unmountDisk "/dev/$target"
sudo dd if="$IMG" of="/dev/r$target" bs=16m status=progress conv=fsync
sync
diskutil eject "/dev/$target"

if [[ "${SKIP_REBOOT:-0}" != 1 ]]; then
  reboot_port="$(resolve_serial_port)"
  stop_rpiboot_mass_storage
  reboot_target_over_serial "$reboot_port"
  echo "Waiting ${REBOOT_SETTLE_SECS}s for target to start booting..."
  sleep "$REBOOT_SETTLE_SECS"
else
  echo "Skipping target reboot because SKIP_REBOOT=1"
fi

if [[ "${SKIP_PROVISION:-0}" != 1 ]]; then
  provision_port="$(wait_for_serial_port "$SERIAL_BOOT_TIMEOUT_SECS")"
  provision_device "$provision_port"
else
  echo "Skipping provisioning because SKIP_PROVISION=1"
fi