#!/bin/bash

set -e
#set -x

# Specify default mount options for certain filesystem types
declare -A DEFAULT_MOUNT_OPTS=(
    ["btrfs"]="defaults,space_cache=v2,discard=async"
    ["vfat"]="rw,relatime,fmask=0022,dmask=0022,codepage=437,iocharset=iso8859-1,shortname=mixed,utf8,errors=remount-ro"
    ["exfat"]="rw,relatime,fmask=0022,dmask=0022,iocharset=utf8,errors=remount-ro"
    ["OTHER"]="defaults"
)

MOUNT_OPTS=""
TARGET=""

[[ -z $SUDO_USER ]] || CURRENT_USER=$SUDO_USER; [[ -z $DOAS_USER ]] || CURRENT_USER=$DOAS_USER
MOUNT_FOLDER="/run/media/${CURRENT_USER}"

ALLOW_UUID="0"
ALLOW_LABEL="0"

print_help () {
    cat /dev/stdin << EOF
$(basename "$0") - Quick Mount

Usage: $(basename "$0") [options] TARGET

Options:
-h    --help              Print this menu
-p    --permissions       Force mount permissions to current user
-u    --uuid              Allow UUIDs as the target
-l    --label             Allow labels as the target
-ro   --read-only         Mounts target as read only
EOF

    exit 0
}

find_target () {
    # If a direct path is specified, always use that first
    if [[ -b "${1}" ]]; then
        TARGET="${1}"
        return
    fi

    if [[ "${ALLOW_UUID}" == "1" ]]; then
        # Search by UUID
        if blkid -o device -t UUID="${1}" 2>/dev/null 1>/dev/null; then
            TARGET=$(blkid -o device -t UUID="${1}")

            if [[ $( echo $TARGET | wc -l ) -gt 1 ]]; then
                echo -e "\e[1;31mMultiple drives contain the UUID \"{$1}\"!""\e[0m" 1>&2
                exit 1
            fi

            return
        fi

        # Search by Part UUID
        if blkid -o device -t PARTUUID="${1}" 2>/dev/null 1>/dev/null; then
            TARGET=$(blkid -o device -t PARTUUID="${1}")

            if [[ $( echo $TARGET | wc -l ) -gt 1 ]]; then
                echo -e "\e[1;31mMultiple drives contain the PARTUUID \"{$1}\"!""\e[0m" 1>&2
                exit 1
            fi

            return
        fi
    fi

    if [[ "${ALLOW_LABEL}" == "1" ]]; then
        # Search by label
        if blkid -o device -t LABEL="${1}" 2>/dev/null 1>/dev/null; then
            TARGET=$(blkid -o device -t LABEL="${1}")

            if [[ $( echo $TARGET | wc -l ) -gt 1 ]]; then
                echo -e "\e[1;31mMultiple drives contain the label \"{$1}\"!""\e[0m" 1>&2
                exit 1
            fi

            return
        fi

        if blkid -o device -t PARTLABEL="${1}" 2>/dev/null 1>/dev/null; then
            TARGET=$(blkid -o device -t PARTLABEL="${1}")

            if [[ $( echo $TARGET | wc -l ) -gt 1 ]]; then
                echo -e "\e[1;31mMultiple drives contain the part label \"{$1}\"!""\e[0m" 1>&2
                exit 1
            fi

            return
        fi
    fi

    # Search by node
    if [[ -b "/dev/${1}" ]]; then
        TARGET="/dev/${1}"
        return
    fi
}

get_mapped () {
    while read -r name type pkname; do
        if [[ "$pkname" == "$dev_name" && "$type" == "crypt" ]]; then
            TARGET=${NAME}
        fi
    done < <(lsblk -lno NAME,TYPE,PKNAME)

    echo "Target \"${TARGET}\" is not unlocked"

    MAPPED_NAME=$(blkid -o value -s UUID "${TARGET}")
    if ! cryptsetup open "${TARGET}" "${MAPPED_NAME}"; then
        echo -e "\e[1;31mCould not unlock target \"${TARGET}\"!""\e[0m" 1>&2
        exit 1
    fi

    TARGET="/dev/mapper/${MAPPED_NAME}"

    trap "cryptsetup close ${TARGET}" EXIT

    [[ ! -z $(blkid -o value -s LABEL "${TARGET}") ]] && LABEL=$(blkid -o value -s LABEL "${TARGET}")

    TYPE=$(blkid -o value -s TYPE "${TARGET}" || true) # Fuck you bash
}

if [[ "$#" -eq 0 ]]; then
    print_help
fi

for arg in "$@"; do
    case "${arg}" in
        "-h" | "--help")
            print_help
        ;;
        "-p" | "--permissions")
            MOUNT_OPTS+=",uid=$UID,gid=$(id -g $UID)"
        ;;
        "-u" | "--uuid")
            ALLOW_UUID="1"
        ;;
        "-l" | "--label")
            ALLOW_LABEL="1"
        ;;
        "-ro" | "--read-only")
            MOUNT_OPTS+=",ro"
        ;;
        *)
            find_target "$arg"
        ;;
    esac
done

if [[ -z $TARGET ]]; then
    echo -e "\e[1;31mA target was not specified!""\e[0m" 1>&2
    exit 1
fi

MAPPED_NAME=""

LABEL=$(blkid -o value -s LABEL "${TARGET}")
[[ -z $LABEL ]] && LABEL=$(basename "${TARGET}")

TYPE=$(blkid -o value -s TYPE "${TARGET}")
[[ "$TYPE" == "crypto_LUKS" ]] && get_mapped
if [[ -z $TYPE ]]; then
    echo -e "\e[1;31mTarget \"${TARGET}\" does not contain a filesystem!""\e[0m" 1>&2
    exit 1
fi

trap - EXIT

if [[ -z "${DEFAULT_MOUNT_OPTS["$TYPE"]}" ]]; then
    MOUNT_OPTS="${DEFAULT_MOUNT_OPTS["OTHER"]}""${MOUNT_OPTS}"
else
    MOUNT_OPTS="${DEFAULT_MOUNT_OPTS["$TYPE"]}""${MOUNT_OPTS}"
fi

[[ ! -d "${MOUNT_FOLDER}" ]] && mkdir "${MOUNT_FOLDER}"
[[ ! -d "${MOUNT_FOLDER}/${LABEL}" ]] && mkdir "${MOUNT_FOLDER}/${LABEL}"

if ! mount -o "${MOUNT_OPTS}" --source "${TARGET}" --target "${MOUNT_FOLDER}/${LABEL}"; then
    echo -e "\e[1;31mFailed to mount \"${TARGET}\"!""\e[0m" 1>&2
    exit 1
fi

(
    set +e

    trap '' SIGHUP SIGTERM

    setsid sh -c '
        while mountpoint "${1}"; do
            sleep 1
        done

        rmdir "${1}"

        [ -b /dev/mapper/"${2}" ] && cryptsetup close /dev/mapper/"${2}"
    ' -- "${MOUNT_FOLDER}/${LABEL}" "${MAPPED_NAME}" >/dev/null 2>&1 &
) &

echo -e "\e[1;32mDevice \"${TARGET}\" successfully mounted at \"${MOUNT_FOLDER}/${LABEL}\"""\e[0m"