#!/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 listing and managing
      components hosted in $(g /srv) in a TechLit system

    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

declare -A SELECTIONS                                                               # Store versions to be deleted
SRV="$BTRFS/@srv"                                                                   # Components directory
COMPONENTS=("admin" "control" "curriculum" "desktop" "help" "recovery" "snapshots") # Components to manage
ANSWER=""                                                                           # Dialog answer placeholder
MOUNTPOINT="$(mktemp -d)"                                                           # USB mountpoint

# 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 "Manage components hosted in $SRV" --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" 6 40
    return 1
  elif ! menu "Select USB Drive" --menu "USB devices" 10 40 20 "${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

  cleanup() { mountpoint -q "$MOUNTPOINT" && run sudo umount -q "$MOUNTPOINT"; }
  return 0
}

# rsync a specific component from USB
sync_component() {
  local component="$1" latest_version silent
  silent="${2:-false}"

  latest_version="$(COMP_DIR="$MOUNTPOINT" tl-comp-latest "$component" 2>/dev/null)"
  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"
    for version in "$MAJOR" "$MINOR" "$PATCH"; do
      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
  for component_dir in $(find "$MOUNTPOINT" -maxdepth 1 -type d 2>/dev/null); do
    comp=$(basename "$component_dir")
    [[ " ${COMPONENTS[*]} " =~ $comp ]] || continue # Filter valid 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
    latest="$(COMP_DIR="$MOUNTPOINT" tl-comp-latest "$component")"
    options+=("$component" "$latest" "off")
  done

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

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

  for component in $selected; do
    sync_component "$component"
  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
  usb_desktop="$(COMP_DIR="$MOUNTPOINT" tl-comp-latest desktop)"

  sync_component admin
  sync_component desktop

  run sudo umount "$MOUNTPOINT"
  say "USB unmounted! You can unplug it now"

  say "Updating Desktop"
  {
    run tl-archive-install admin
    run tl-img-unpack "$BTRFS" "$usb_desktop"
    run tl-img-restore
  }
  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
}

# Deletes all selected components in SELECTIONS
apply_selections() {
  local pre_disk_usage post_disk_usage options

  pre_disk_usage=$(sudo df -h "$BTRFS" --output=used -h | tail -n1)

  if [ ${#SELECTIONS[@]} -gt 0 ]; then
    for comp in "${!SELECTIONS[@]}"; do
      options+="\n$comp: ${SELECTIONS["$comp"]}\n"
    done

    if ! menu "Summary of Selections" --yesno \
      "\Zb\Z1 Warning!!!\Zn You have selected to delete these \n
       \Zb\Z3 $options\n \Zn
       \Zb\Z1 This is irreversible. Proceed?" 20 50; then
      return
    fi
  else
    menu "\Z1Nothing selected" --pause "" 7 20 3
    return
  fi

  for component in "${!SELECTIONS[@]}"; do
    for version in ${SELECTIONS["$component"]}; do
      ver="$SRV/$component/$version"
      [ -z "$ver" ] && return

      # If the component is snapshot, delete the subvolumes first
      if [ -d "$ver".d/@root ] || [ -d "$ver".d/@guest ]; then
        run sudo btrfs subvol delete "$ver".d/@{root,guest} || continue
      fi
      run sudo rm -rf "$ver".*
    done
  done

  run sudo sync
  post_disk_usage="$(sudo df -h "$BTRFS" --output=used -h | tail -n1)"

  menu "Components deleted successfully!" \
    --pause "Initial Disk Usage: $pre_disk_usage\n
Final Disk Usage:   $post_disk_usage" 10 60 5

  SELECTIONS=()
}

purge_components() {
  if ! menu "Purge Components" --yesno \
    "\ZbDelete all temporary files?" 6 60; then
    return
  fi

  # Incomplete unpacks
  for dir in $(sudo find "$SRV"/snapshots -maxdepth 1 -type d -name 'tmp-*'); do
    if [ -d "$dir/@root" ]; then
      run sudo btrfs subvol delete "$dir/@{root,guest}" || :
    fi
    run sudo rm -rf "$dir"
  done

  # Downloaded using aria2
  run find "$SRV" \
    \( -path "$SRV"/snapshots -o -path "$SRV"/secure -o -path "$SRV"/public \) -prune -o \
    -type f -name '*.aria2' -print | while read -r tmp_file; do

    file="${tmp_file%.aria2}"
    run rm -f "$file"*
  done || :

  # incomplete rsync downloads
  for comp in "${COMPONENTS[@]}"; do
    [ -d "$SRV/$comp/.rsync" ] || continue
    run sudo rm -rf "$SRV/$comp"/.rsync/*
  done
}

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

# Lists a component and its versions as options and save them to SELECTIONS
sub_menu() {
  local component version_name usage options=() selected_versions=()
  component="$1"

  if [ -n "${SELECTIONS[$component]}" ]; then
    selected_versions=("${SELECTIONS["$component"]}")
  fi

  for version in "$SRV/$component"/*.{tar.gz,iso,d}; do
    [ -r "$version" ] || continue

    version_name="$(basename "$(echo "$version" | sed -E 's/\.(tar\.gz|iso|d)$//')")"
    usage="$(sudo du -sxh "$version" | awk '{print $1}')"

    [[ " ${selected_versions[*]} " == *" $version_name "* ]] \
      && options+=("$version_name" "$usage" on)
    options+=("$version_name" "$usage" off)
  done

  if [ ${#options[@]} -eq 0 ]; then
    menu "Selected" --msgbox "\Zb\Z1404:\Zn No component found" 6 30
    main_menu
    return
  fi

  menu "Versions Menu" \
    --checklist "Select versions to delete:" \
    20 40 6 "${options[@]}" || :

  if [ -n "$ANSWER" ]; then
    SELECTIONS["$component"]="$ANSWER"
  fi
  ANSWER=""

  main_menu
}

# The main menu: Loops over components check if it exists and display them as options
# Then run sub_menu with the selected component
main_menu() {
  local disk_usage version options=() latest usage
  disk_usage=$(sudo df -h "$BTRFS" --output=size,used,avail,pcent \
    | tail -n1 | awk '{printf "Total: %s | Available: %s | Used: %s", $1, $3, $4}')

  while true; do
    options=() # Resetting here to avoid multiple lists
    for component in "${COMPONENTS[@]}"; do
      [ -d "$SRV/$component" ] || continue
      version="$(tl-comp-latest "$component" 2>/dev/null)"
      version="${version:-None}"
      latest="Latest \Zb\Z6 $(printf "%-10s" "$version")\Zn"
      usage="Disk \Zb\Z3$(sudo du -sxh "$(readlink -f "$SRV/$component")" \
        2>/dev/null | awk '{printf "%-5s", $1}')\Zn"
      options+=("$component" "$usage  $latest")
    done

    options+=("Apply" "\Z1Delete all selected components")
    options+=(" " "") # A dummy separator
    options+=("USB" "Update components (from USB)")
    options+=("Update" "Update Desktop (Using USB)")
    options+=("Purge" "Temporary files/incomplete downloads")
    options+=("Quit" "Exit without changes")

    menu "Select component or an option" \
      --menu "\Zb\Z4Disk Usage\Zn ⇒ $disk_usage" 21 60 10 "${options[@]}" || :

    [ -z "$ANSWER" ] && return

    case "$ANSWER" in
      " ") ;;
      Apply) apply_selections ;;
      USB) copy_from_usb ;;
      Update) update_system ;;
      Purge) purge_components ;;
      Quit) quit_program ;;
      *) sub_menu "$ANSWER" ;;
    esac
    ANSWER=""
  done
}

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