#!/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="${MOUNT_OPTS:-""}"
TARGET="${TARGET:-""}"

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

ALLOW_UUID="${ALLOW_UUID:-0}"
ALLOW_LABEL="${ALLOW_LABEL:-0}"
LABEL_BY="${LABEL_BY:-LABEL}"

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
-d    --dry-run           Does nothing
-U    --label-uuid        Use UUID instead of drive label for mountpoint name
-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"

    if [[ -z $DRY_RUN ]]; then
        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_BY}" "${TARGET}") ]] && LABEL=$(blkid -o value -s "${LABEL_BY}" "${TARGET}")

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

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"
        ;;
        "-d" | "--dry-run")
            DRY_RUN="1"
        ;;
        "-U" | "--label-uuid")
            LABEL_BY="UUID"
        ;;
        "-ro" | "--read-only")
            MOUNT_OPTS+=",ro"
        ;;
        *)
            [[ -z $TARGET ]] && 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

TARGET_TYPE=$(lsblk -ndo TYPE "${TARGET}")
if [[ ( "$TARGET_TYPE" == "disk" || "$TARGET_TYPE" == "loop" ) && -z "$(blkid -o value -s TYPE "${TARGET}")" ]]; then
    lsblk --list --paths --noheadings --output NAME,TYPE "${TARGET}" | while read name type; do
        [[ "$type" = "part" ]] && ( TARGET="$name" "$0" $@ ) </dev/tty
    done
    exit 0
fi

MAPPED_NAME=""

LABEL=$(blkid -o value -s "${LABEL_BY}" "${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}" && -z $DRY_RUN ]] && mkdir "${MOUNT_FOLDER}"
[[ ! -d "${MOUNT_FOLDER}/${LABEL}" && -z $DRY_RUN ]] && mkdir "${MOUNT_FOLDER}/${LABEL}"


if [[ -z $DRY_RUN ]]; then
    if \
        mountpoint "${MOUNT_FOLDER}/${LABEL}" 2>&1 &>/dev/null || \
        findmnt -n -o SOURCE "${TARGET}" 2>&1 &>/dev/null || \
        ! 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 &
    ) &
fi

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