#!/usr/bin/env bash
# vim:et:ts=2:sts=2:sw=2

# Get absolute repository root (especially when symlinked)
ROOT="$(realpath "$(dirname \"$(dirname \"$(readlink -f \"${BASH_SOURCE[0]}\")\")\")")"

# Load accompanying bash library
source "$ROOT/lib/lib.bash"

BTRFS="${1:-/var/btrfs}"

if [ -n "$HELP" ] || [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
  say "
  USAGE:
    $(g "${BASH_SOURCE##*/}")
  DESCRIPTION:
    $(g "${BASH_SOURCE##*/}") Provides a TUI for updating components
    from a USB drive.
  OPTIONS:
    -h, / --help      Show this help message and exit
  "
  exit 1
fi

# Exit now unless run as admin
assert-is-admin

say "This script requires root access"
sudo echo thank you

COMPONENTS=("admin" "control" "curriculum" "desktop" "help" "recovery" "snapshots") # Versioned components to manage
ANSWER=""                                                                           # Dialog answer placeholder
MOUNTPOINT="$(mktemp -d)"                                                           # USB mountpoint
LOCKFILE="/tmp/tl-comp-update.lock"                                                 # Lockfile

if [ -e "$LOCKFILE" ]; then
  err "Another instance is already running. Exiting."
  exit 1
fi

cleanup() {
  mountpoint -q "$MOUNTPOINT" && run sudo umount -q "$MOUNTPOINT"
  rm -f "$LOCKFILE"
}

# Helper function for displaying dialog menus
# USAGE: menu "Title" --yesno "Message" 10 60
#        menu "Title" --checklist "Message" 20 40 6 "Option 1" "Description 1" "Option 2" "Description 2"
menu() {
  ANSWER="$(
    dialog \
      --clear --no-collapse --output-fd 1 --colors \
      --backtitle "Update components from USB" --title "$@"
  )"
}

# Select and mount a USB device
# USAGE: select_and_mount_usb
select_and_mount_usb() {
  local options=()

  for block in $(lsblk -o NAME,TRAN | grep usb | awk '{print $1}'); do
    options+=("/dev/$block" "USB Device")
  done

  if [ ${#options[@]} -eq 0 ]; then
    menu "No USB Drive" --msgbox "No USB drive found check the port" 6 40
    return 1
  elif ! menu "Select USB Drive" --menu "USB devices" 10 50 2 "${options[@]}"; then
    return 1
  else
    usb="$ANSWER"
  fi
  if ! sudo mount "$usb"4 "$MOUNTPOINT"; then
    menu "Error" --msgbox \
      "Failed to mount the selected USB drive. Please check the drive and try again" 6 60
    return 1
  fi

  return 0
}

# rsync a specific component from USB
sync_component() {
  local component="$1"
  local force_copy="${2:-false}"
  local latest_version
  local SRV="$BTRFS/@srv"

  if [ "$component" = "desktop" ] && is-mini "$BTRFS"; then
    if [ -d "$MOUNTPOINT/desktop-mini" ]; then
      component="desktop-mini"
    fi
  fi

  if [ "$component" = "curriculum" ]; then
    return
  else
    latest_version="$(COMP_DIR="$MOUNTPOINT" "$ROOT/bin/tl-comp-latest" "$component" 2>/dev/null)"
  fi
  say "Copying $component $latest_version"

  if [[ "$component" =~ ^desktop ]]; then
    PATCH="$latest_version"
    MINOR="$(echo "$PATCH" | cut -d. -f-2).0"
    MAJOR="$(echo "$PATCH" | cut -d. -f-1).0.0"
    local check_ext="d"
    is-mini "$BTRFS" && check_ext="mini.d"

    for version in "$MAJOR" "$MINOR" "$PATCH"; do
      if [ "$force_copy" = "false" ] && [ -d "$BTRFS/@srv/snapshots/v$version.$check_ext/@root" ]; then
        say "Snapshot for $component v$version already exists. Skipping copy."
        continue # Skip if snap is already there AND not forcing copy
      fi
      run rsync -avP "$MOUNTPOINT/$component/v$version"* "$SRV/$component/"
    done
  else
    run rsync -avP "$MOUNTPOINT/$component/v$latest_version"* "$SRV/$component/"
  fi
}

copy_from_usb() {
  ! select_and_mount_usb && return

  local components=() options=() selected latest info
  local SRV="/srv"
  for component_dir in $(find "$MOUNTPOINT" -maxdepth 1 -type d 2>/dev/null); do
    comp=$(basename "$component_dir")
    [[ " ${COMPONENTS[*]} " =~ $comp ]] || continue # Filter out invalid components
    components+=("$comp")
  done

  if [ ${#components[@]} -eq 0 ]; then
    menu "No Components Found" --msgbox "No valid components found on the USB drive." 6 40
    return
  fi

  for component in "${components[@]}"; do
    if [ "$component" = "curriculum" ]; then
      info="$(sudo du -sh "$MOUNTPOINT/$component" | awk '{print $1}')"
    else
      info="$(COMP_DIR="$MOUNTPOINT" "$ROOT/bin/tl-comp-latest" "$component" 2>/dev/null || echo "Unknown")"
    fi
    options+=("$component" "$info" "off")
  done

  # Select components to sync
  if ! menu "Select Components" --checklist "Choose components to sync:" 20 50 10 "${options[@]}"; then
    cleanup
    return
  fi

  [ -z "$ANSWER" ] && return
  selected="$ANSWER"

  for component in $selected; do
    if [ "$component" = "desktop" ]; then
      sync_component "$component" true
    else
      sync_component "$component"
    fi
  done

  menu "Sync Complete" --msgbox "Success syncing: \Zb\Z2$selected" 6 40
  cleanup
}

update_system() {
  if [ -r "/tmp/tl-did-restore" ]; then
    err "Reboot to proceed!"
    exit 1
  fi

  ! select_and_mount_usb && return

  local usb_comp="desktop"
  if is-mini "$BTRFS" && [ -d "$MOUNTPOINT/desktop-mini" ]; then
    usb_comp="desktop-mini"
  fi
  local usb_desktop="$(COMP_DIR="$MOUNTPOINT" "$ROOT/bin/tl-comp-latest" "$usb_comp")"

  DISK_SIZE_GB=$(df --output=size -B1G "$BTRFS" | tail -n1)

  # This enables us to use a USB drive for staging in <32GB storage devices
  if [[ "$DISK_SIZE_GB" -lt 32 ]]; then
    say "Low storage device detected. Using USB for staging."
    STAGE="$MOUNTPOINT/stage"
    mkdir -p "$STAGE"
    ARCH="$MOUNTPOINT/$usb_comp"
    STAGE="$STAGE" ARCH="$ARCH" run "$ROOT/bin/tl-img-unpack" "$BTRFS" "$usb_desktop"

    run sudo umount "$MOUNTPOINT"
    menu "USB Ready" --msgbox "USB unmounted! You can unplug it now." 6 40
  else


    sync_component desktop "false"

    # USB is no longer needed after sync! Unmount early.
    run sudo umount "$MOUNTPOINT"
    menu "USB Ready" --msgbox "USB unmounted! You can unplug it now." 6 40

    run "$ROOT/bin/tl-img-unpack" "$BTRFS" "$usb_desktop"
  fi

  say "Updating Desktop"
  {
    run "$ROOT/bin/tl-archive-install" admin
    run "$ROOT/bin/tl-img-restore" "$BTRFS" "$usb_desktop"
  }
  while true; do
    echo "$(y Reboot now?) $(c '[Y/n]')"
    read -n1 -s -r x

    case "${x:-y}" in
      y | Y)
        run sudo reboot
        break
        ;;
      n | N) break ;;
    esac
  done

  exit 0
}

quit_program() {
  if ! menu "Quit" --yesno "quit?" 6 60; then
    return
  fi
  exit 0
}

# The main menu
main_menu() {
  touch "$LOCKFILE" || {
    err "Failed to create lock file. Exiting."
    exit 1
  }

  while true; do
    options=()
    options+=("USB" "Update components (from USB)")
    options+=("Update" "Update Desktop (Using USB)")
    options+=("Quit" "Exit without changes")

    menu "Select an option" \
      --menu "" 10 65 5 "${options[@]}" || :

    [ -z "$ANSWER" ] && return

    case "$ANSWER" in
      USB) copy_from_usb ;;
      Update) update_system ;;
      Quit) quit_program ;;
      *) ;;
    esac
    ANSWER=""
  done
}

cleanup
[[ "$0" = "${BASH_SOURCE[0]}" ]] && main_menu
