#!/bin/bash

# This script is meant to support common media managing utilities.
#
# Parameters:
#  - list           -- list media
#  - mount <device> -- mount a device
#  - unmount <path> -- unmount given mount point
#
# Supported utilities (only one is used, listed higher to lower priority):
#  - udevil: udisks-like command-line tool with no D-Bus or daemon
#    requirements: `udevil` and `devmon` commands (both are part of udevil)
#  - udisks: first vesion of disk management D-Bus daemon
#    requirements: `udisks` and `umount` commands
#  - udisks2: second vesion of disk management D-Bus daemon
#    requirements: `udisksctl` and either `python2` or `python3` with "dbus"
#                  module
#
# Support for additional protocols (used in addition to standard drives, but
# support depends on installed dependencies):
#  - MTP: /dev/libmtp-* devices
#    requirements: requires `simple-mtpfs` command to obtain device's name or
#                  mount/unmount it
#    limitations: can list and unmount only if this script was used for mounting
#                 because source field of a FUSE mount is mounter's name rather
#                 than the mounted device

function usage_error() {
    echo "Usage: vifm-media list | mount <device> | unmount <path>"
    exit 1
}

function list() {
    while read -r dev; do
        if [ "${dev:0:10}" = /dev/disk/ ]; then
            continue
        fi

        local out=$(info "$dev")

        if grep -qe '^\s*partition table:\s*$' <<< "$out"; then
            # can't mount the whole drive
            continue
        fi

        if grep -qe '^\s*removable:\s*0\s*$' <<< "$out" &&
           grep -qe '^\s*system internal:\s*1\s*$' <<< "$out"; then
            # skip drives which aren't detachable
            continue
        fi

        if grep -qe '^\s*has media:\s*0\s*$' <<< "$out"; then
            continue
        fi

        echo device="$dev"

        local label=$(grep -m1 '^\s*label:.*$' <<< "$out" | sed 's/^[^:]*:\s*//')
        echo "label=$label"

        local paths=$(grep -m1 '^\s*mount paths:.*$' <<< "$out" | sed 's/^[^:]*:\s*//')
        IFS=',' paths=( $paths )
        for path in "${paths[@]}"; do
            echo "mount-point=$path"
        done

        echo
    done <<< "$("$monitor" --enumerate-device-files)"
}

if type -P devmon udevil &>/dev/null; then
    monitor=devmon

    function info() {
        udevil info "$@"
    }
    function mount() {
        devmon --mount "$1"
    }
    function unmount() {
        devmon --unmount "$1"
    }
elif type -P udisks umount &>/dev/null; then
    monitor=udisks

    function info() {
        udisks --show-info "$@"
    }
    function mount() {
        udisks --mount "$1"
    }
    function unmount() {
        umount "$1"
    }
elif python2 -c 'import dbus' &>/dev/null ||
     ${python3:=python3} -c 'import dbus' &>/dev/null &&
     type -P udisksctl &>/dev/null; then
    python=${python3:-python2}
    function udisks2() {
        $python - "$@" <<"EOF"
import dbus
import sys

def decode(array):
    return bytearray(array).replace(b'\x00', b'').decode('utf-8')

def devs():
    bus = dbus.SystemBus()
    ud_manager_obj = bus.get_object('org.freedesktop.UDisks2',
                                    '/org/freedesktop/UDisks2')
    om = dbus.Interface(ud_manager_obj, 'org.freedesktop.DBus.ObjectManager')
    for v in om.GetManagedObjects().values():
        drive_info = v.get('org.freedesktop.UDisks2.Block', {})
        part_table = v.get('org.freedesktop.UDisks2.PartitionTable')
        usage = drive_info.get('IdUsage')
        system = drive_info.get('HintSystem')
        ro = drive_info.get('ReadOnly')
        if usage == "filesystem" and \
           not system and \
           not ro and \
           part_table is None:
            device = drive_info.get('Device')
            if device is not None:
                device = decode(device)
                fs = v.get('org.freedesktop.UDisks2.Filesystem', {})
                yield (device, drive_info.get('IdLabel'), fs)

def list_usb():
    for (device, label, fs) in devs():
        print('device=%s' % device)
        print('label=%s' % label)
        for point in fs.get('MountPoints', {}):
            point = decode(point)
            print("mount-point=%s" % point)

def path_to_dev(path):
    import os
    path = os.path.normpath(path)
    for (device, label, fs) in devs():
        for point in fs.get('MountPoints', {}):
            point = os.path.normpath(decode(point))
            if point == path:
                return device
    return None

def mount(device):
    bus = dbus.SystemBus()
    obj = bus.get_object('org.freedesktop.UDisks2',
                         '/org/freedesktop/UDisks2/block_devices%s'%device[4:])
    print(obj.Mount({}, dbus_interface='org.freedesktop.UDisks2.Filesystem'))

def unmount(path):
    device = path_to_dev(path)
    if device is None:
        sys.exit(1)

    bus = dbus.SystemBus()
    obj = bus.get_object('org.freedesktop.UDisks2',
                         '/org/freedesktop/UDisks2/block_devices%s'%device[4:])
    obj.Unmount({}, dbus_interface='org.freedesktop.UDisks2.Filesystem')

if len(sys.argv) < 2 or sys.argv[1] == 'list':
    list_usb()
elif len(sys.argv) == 3 and sys.argv[1] == 'mount':
    mount(sys.argv[2])
elif len(sys.argv) == 3 and sys.argv[1] == 'unmount':
    unmount(sys.argv[2])
EOF
    }
    function list() {
        udisks2 list
    }
    function mount() {
        udisks2 mount "$1"
    }
    function unmount() {
        udisks2 unmount "$1"
    }
else
    echo "Neither of the supported backends were found" \
         "(udevil, udisks, udisks2)" >&2
    exit 1
fi

function have_simple_mtpfs() {
    type simple-mtpfs > /dev/null 2>&1
}

function make_mtp_mount_path() {
    local dev=$1

    echo "/tmp/vifm-media$(echo -n "$dev" | tr -c 'a-zA-Z0-9-' _)"
}

function is_mtp_mount_path() {
    local path=$1

    case "$path" in
        /tmp/vifm-media*) return 0 ;;
        *)                return 1 ;;
    esac
}

function list_mtp() {
    for dev in /dev/libmtp-*; do
        echo device="$dev"

        if have_simple_mtpfs; then
            local label=$(simple-mtpfs --list-devices "$dev" | cut -f2- -d' ')
            echo "label=$label"
        fi

        echo -n /dev/libmtp-1-9 | tr -c 'a-zA-Z0-9-' _

        local mount_path=$(make_mtp_mount_path "$dev")
        if findmnt "$mount_path"; then
            echo "mount-point=$mount_path"
        fi
    done
}

function is_mtp() {
    local dev=$1

    case "$dev" in
        /dev/libmtp-*) return 0 ;;
        *)             return 1 ;;
    esac
}

function mount_mtp() {
    local dev=$1

    if ! have_simple_mtpfs; then
        echo "vifm-media: need simple-mtpfs command to mount over MTP" 1>&2
        exit 1
    fi

    local mount_path=$(make_mtp_mount_path "$dev")
    if [ -e "$mount_path" ]; then
        if ! rmdir "$mount_path"; then
            exit 1
        fi
    fi

    if ! ( umask 077 && mkdir "$mount_path" ); then
        exit 1
    fi

    if ! simple-mtpfs "$dev" "$mount_path"; then
        rmdir "$mount_path"
        exit 1
    fi
}

function unmount_mtp() {
    local path=$1

    umount "$path" && rmdir "$path"
}

case "$1" in
    list)
        if [ $# -ne 1 ]; then
            usage_error
        fi
        list
        list_mtp
        ;;
    mount)
        if [ $# -ne 2 ]; then
            usage_error
        fi

        if is_mtp "$2"; then
            mount_mtp "$2"
        else
            mount "$2"
        fi
        ;;
    unmount)
        if [ $# -ne 2 ]; then
            usage_error
        fi

        if is_mtp_mount_path "$2"; then
            unmount_mtp "$2"
        else
            unmount "$2"
        fi
        ;;

    *)
        usage_error
        ;;
esac