#!/usr/bin/env bash

# Description: Toggle mount status of a device using pmount
#              If the device is not mounted, it will be mounted.
#              If the device is mounted, it will be unmounted and powered down.
#
# Dependencies: lsblk, pmount (optional), udisks2
#
# Usage: Runs `lsblk` on 'l', unmounts all USB HDDs on 'u', exits on 'Return`.
#
# Notes:
#   - The script uses Linux-specific lsblk to list block devices. Alternatives:
#       macOS: "diskutil list"
#       BSD: "geom disk list"
#   - The script uses udisksctl (from udisks2) to power down devices. This is also Linux-specific.
#     Users on non-Linux platforms can comment it and use an alternative to power-down disks.
#
# Shell: Bash
# Author: Arun Prakash Jana

prompt="device name [e.g. sdXn] ('l'ist, 'u'nmount all USB, 'q'uit): "
previous_input=""

restore_previous_input() {
    if [ -n "$previous_input" ]; then
        READLINE_LINE="$previous_input"
        READLINE_POINT=${#READLINE_LINE}
    fi
}

clear_input() {
    READLINE_LINE=""
    READLINE_POINT=0
}

bind -x '"\e[A":restore_previous_input'
bind -x '"\e[B":clear_input'

lsblk

printf "\nEnsure you aren't still in the mounted device.\n"
read -r -e -p "$prompt" dev

while [ -n "$dev" ]; do
    previous_input="$dev"
    if [ "$dev" = "l" ]; then
        lsblk
    elif [ "$dev" = "q" ]; then
        exit
    elif [ "$dev" = "u" ]; then
        usb_devs=$(lsblk -ndo NAME,TRAN | awk '$2=="usb"{print $1}')
        if [ -z "$usb_devs" ]; then
            echo "No USB block devices found."
        else
            for disk in $usb_devs; do
                # Get only real partitions (e.g. sda1), skip device-mapper entries like luks-*
                parts=$(lsblk -nlpo NAME,TYPE "/dev/$disk" | awk '$2=="part"{print $1}')
                if [ -z "$parts" ]; then
                    echo "/dev/$disk has no partitions, skipping."
                    continue
                fi
                mounted=0
                for part in $parts; do
                    mnt=$(lsblk -n "$part" -o MOUNTPOINT | sed '/^$/d')
                    # Check if this partition or any child (e.g. LUKS) is mounted
                    child_mnt=$(lsblk -nro MOUNTPOINT "$part" | sed '/^$/d')
                    if [ -z "$mnt" ] && [ -z "$child_mnt" ]; then
                        continue
                    fi
                    mounted=1
                    if type pumount >/dev/null 2>&1; then
                        sync "$child_mnt"
                        pumount "$part"
                        prc=$?
                    elif lsblk -n "$part" -o FSTYPE | grep -q "crypto_LUKS"; then
                        dev_map=$(udisksctl info -b "$part" | grep "CleartextDevice" | grep -o "dm_2d[[:digit:]]*" | sed "s/_2d/-/")
                        sync "$child_mnt"
                        udisksctl unmount -b "/dev/$dev_map" --no-user-interaction >/dev/null
                        prc=$?
                        udisksctl lock -b "$part" --no-user-interaction >/dev/null
                    else
                        sync "$mnt"
                        udisksctl unmount -b "$part" --no-user-interaction >/dev/null
                        prc=$?
                    fi
                    if [ $prc -eq 0 ]; then
                        echo "$part unmounted."
                    else
                        echo "Failed to unmount $part."
                    fi
                done
                if [ $mounted -eq 0 ]; then
                    echo "/dev/$disk has no mounted partitions, skipping."
                    continue
                fi
                if udisksctl power-off -b "/dev/$disk" --no-user-interaction 2>/dev/null; then
                    echo "/dev/$disk ejected."
                fi
            done
        fi
    else
        # LUKS volumes mounted with udisksctl appear differently than with pmount
        if grep -qs "$dev " /proc/mounts || [ -n "$(lsblk -n "/dev/$dev" -o MOUNTPOINT)" ]; then
            sync "$(lsblk -n "/dev/$dev" -o MOUNTPOINT | sed "/^$/d")"
            if type pumount >/dev/null 2>&1; then
                pumount "/dev/$dev"
                exit_code="$?"
            else
                # Unlike pmount, udisksctl does not transparently handle LUKS volumes
                # We need to manually get the unlocked device, unmount it, and then lock the volume
                if lsblk -n "/dev/$dev" -o FSTYPE | grep "crypto_LUKS" >/dev/null; then
                    dev_map="$(udisksctl info -b /dev/"$dev" | grep "CleartextDevice" | grep -o "dm_2d[[:digit:]]*" | sed "s/_2d/-/")"
                    udisksctl unmount -b "/dev/$dev_map" --no-user-interaction >/dev/null
                    exit_code="$?"
                    udisksctl lock -b "/dev/$dev" --no-user-interaction >/dev/null
                else
                    udisksctl unmount -b "/dev/$dev" --no-user-interaction >/dev/null
                    exit_code="$?"
                fi
            fi
            if [ $exit_code -eq 0 ]; then
                echo "/dev/$dev unmounted."
                if udisksctl power-off -b "/dev/$dev" --no-user-interaction; then
                    echo "/dev/$dev ejected."
                fi
            fi
        elif [ -b "/dev/$dev" ]; then
            if type pmount >/dev/null 2>&1; then
                pmount "/dev/$dev"
                exit_code="$?"
            else
                # Unlike pmount, udisksctl does not transparently handle LUKS volumes
                # We need to manually get the unlocked device and mount that instead of the volume itself
                if [ "$(lsblk "/dev/$dev" -n -o FSTYPE)" = "crypto_LUKS" ]; then
                    dev_map=$(udisksctl unlock -b "/dev/$dev" --no-user-interaction | grep -o "dm-[[:digit:]]*")
                    udisksctl mount -b "/dev/$dev_map" --no-user-interaction >/dev/null
                    exit_code="$?"
                else
                    udisksctl mount -b "/dev/$dev" --no-user-interaction >/dev/null
                    exit_code="$?"
                fi
            fi
            if [ $exit_code -eq 0 ]; then
                sleep 1
                echo "/dev/$dev mounted to $(lsblk -n "/dev/$dev" -o MOUNTPOINT | sed "/^$/d")."
            fi
        else
            echo "/dev/$dev does not exist or is not a block device."
        fi
    fi

    echo
    read -r -e -p "$prompt" dev
done