#!/usr/bin/env bash
# Waybar USB module - automount, status, open in yazi, eject
# Usage:
#   usb.sh          -> JSON status for waybar poll (auto-mounts unmounted USBs)
#   usb.sh open     -> open first mounted USB in yazi (via kitty)
#   usb.sh eject    -> unmount + power-off all USB devices
set -euo pipefail

# --- helpers -------------------------------------------------------------------

get_usb_devices() {
    # Returns JSON array of USB block devices via lsblk.
    # Fields: NAME, RM (removable), SIZE, LABEL, MOUNTPOINT, FSTYPE, TRAN, MODEL
    lsblk -Jbpno NAME,RM,SIZE,LABEL,MOUNTPOINT,FSTYPE,TRAN,MODEL 2>/dev/null \
        | jq -c '[.blockdevices[] | select(.tran == "usb")]'
}

auto_mount_unmounted() {
    local devices="$1"
    # Auto-mount any USB partition that isn't mounted yet
    local unmounted
    unmounted=$(echo "$devices" | jq -r '.[] | select(.mountpoint == null or .mountpoint == "") | .name')
    for dev in $unmounted; do
        udisksctl mount -b "$dev" --no-user-interaction 2>/dev/null || true
    done
}

# --- actions -------------------------------------------------------------------

do_open() {
    # Open the first mounted USB in yazi via kitty
    local devices
    devices=$(get_usb_devices)

    local mountpoint
    mountpoint=$(echo "$devices" | jq -r '[.[] | select(.mountpoint != null and .mountpoint != "")] | first | .mountpoint // empty')

    if [ -z "$mountpoint" ]; then
        # Try children (partitions)
        mountpoint=$(lsblk -Jbpno NAME,RM,SIZE,LABEL,MOUNTPOINT,FSTYPE,TRAN,MODEL 2>/dev/null \
            | jq -r '[.blockdevices[] | select(.tran == "usb") | .children[]? | select(.mountpoint != null and .mountpoint != "")] | first | .mountpoint // empty')
    fi

    if [ -n "$mountpoint" ]; then
        kitty -e yazi "$mountpoint" &
        disown
    else
        notify-send "USB" "No mounted USB device found" -i dialog-warning
    fi
}

do_eject() {
    # Unmount all USB partitions, then power-off the parent device
    local devices
    devices=$(get_usb_devices)

    local mounted
    mounted=$(echo "$devices" | jq -r '.[] | select(.mountpoint != null and .mountpoint != "") | .name')

    # Also check children (partitions of whole disks)
    local child_mounted
    child_mounted=$(lsblk -Jbpno NAME,MOUNTPOINT,TRAN 2>/dev/null \
        | jq -r '[.blockdevices[] | select(.tran == "usb") | .children[]? | select(.mountpoint != null and .mountpoint != "")] | .[].name')

    for dev in $mounted $child_mounted; do
        udisksctl unmount -b "$dev" --no-user-interaction 2>/dev/null || true
    done

    # Power-off the whole USB devices
    local parent_devs
    parent_devs=$(echo "$devices" | jq -r '.[].name')
    for dev in $parent_devs; do
        udisksctl power-off -b "$dev" --no-user-interaction 2>/dev/null || true
    done

    notify-send "USB" "All USB devices ejected" -i media-eject
}

do_status() {
    local devices
    devices=$(get_usb_devices)

    # Auto-mount unmounted partitions
    auto_mount_unmounted "$devices"
    # Re-read after automount
    devices=$(get_usb_devices)

    # Count mounted USBs (check both top-level and children)
    local count=0
    local tooltip_lines=""

    # Top-level devices
    while IFS= read -r line; do
        if [ -z "$line" ] || [ "$line" = "null" ]; then continue; fi
        local name size label mount fstype model
        name=$(echo "$line" | jq -r '.name')
        size=$(echo "$line" | jq -r '.size')
        label=$(echo "$line" | jq -r '.label // "N/A"')
        mount=$(echo "$line" | jq -r '.mountpoint // "not mounted"')
        fstype=$(echo "$line" | jq -r '.fstype // "?"')
        model=$(echo "$line" | jq -r '.model // "USB" | gsub("^\\s+|\\s+$"; "")')

        if [ "$mount" != "not mounted" ]; then
            count=$((count + 1))
        fi
        tooltip_lines="${tooltip_lines}${name} ${model} ${size} [${fstype}]\nLabel: ${label}\nMount: ${mount}\n"
    done < <(echo "$devices" | jq -c '.[]')

    # Children (partitions)
    while IFS= read -r line; do
        if [ -z "$line" ] || [ "$line" = "null" ]; then continue; fi
        local name size label mount fstype
        name=$(echo "$line" | jq -r '.name')
        size=$(echo "$line" | jq -r '.size')
        label=$(echo "$line" | jq -r '.label // "N/A"')
        mount=$(echo "$line" | jq -r '.mountpoint // "not mounted"')
        fstype=$(echo "$line" | jq -r '.fstype // "?"')

        if [ "$mount" != "not mounted" ]; then
            count=$((count + 1))
        fi
        tooltip_lines="${tooltip_lines}  └ ${name} ${size} [${fstype}]\n    Label: ${label}\n    Mount: ${mount}\n"
    done < <(lsblk -Jbpno NAME,RM,SIZE,LABEL,MOUNTPOINT,FSTYPE,TRAN,MODEL 2>/dev/null \
        | jq -c '[.blockdevices[] | select(.tran == "usb") | .children[]?] | .[]')

    # Determine alt class
    local alt="usb-none"
    if [ "$count" -gt 0 ]; then
        alt="usb-mounted"
    elif [ -n "$tooltip_lines" ]; then
        alt="usb-unmounted"
    fi

    # Build JSON output
    local text icon
    if [ "$count" -gt 0 ]; then
        icon="󰋊"
        text="${icon} ${count}"
    elif [ -n "$tooltip_lines" ]; then
        icon="󰋊"
        text="${icon} !"
    else
        text=""
    fi

    # Escape for JSON
    local tooltip_json
    tooltip_json=$(echo -e "$tooltip_lines" | head -20 | jq -Rs '.')

    if [ -n "$text" ]; then
        printf '{"text":%s,"tooltip":%s,"alt":%s}\n' \
            "$(jq -n --arg t "$text" '$t')" \
            "$tooltip_json" \
            "$(jq -n --arg a "$alt" '$a')"
    else
        echo '{"text":"","tooltip":"","alt":"usb-none"}'
    fi
}

# --- main ----------------------------------------------------------------------

case "${1:-}" in
    open)    do_open ;;
    eject)   do_eject ;;
    *)       do_status ;;
esac