#!/bin/bash
#
# Interactive USB Auto-mount Setup
# Scans for connected USB drives and lets you select which ones to auto-mount
#

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

echo "=========================================="
echo "Interactive USB Auto-mount Setup"
echo "=========================================="
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}This script needs sudo privileges${NC}"
    echo "Please run with: sudo ./setup_usb_interactive.sh"
    exit 1
fi

# Get the real user
REAL_USER=${SUDO_USER:-$USER}
REAL_UID=$(id -u $REAL_USER)
REAL_GID=$(id -g $REAL_USER)

echo "Setting up for user: $REAL_USER (UID:$REAL_UID)"
echo ""

# Step 1: Stop existing systemd units
echo "Step 1: Stopping existing systemd units..."
systemctl stop media-usb*.automount 2>/dev/null || true
systemctl stop media-usb*.mount 2>/dev/null || true
systemctl disable media-usb*.automount 2>/dev/null || true
echo -e "${GREEN}✓ Stopped existing units${NC}"
echo ""

# Step 2: Unmount USB drives
echo "Step 2: Unmounting existing USB mounts..."
for mount in /media/usb*; do
    if mountpoint -q "$mount" 2>/dev/null; then
        umount "$mount" 2>/dev/null || true
    fi
done
echo -e "${GREEN}✓ Unmounted drives${NC}"
echo ""

# Step 3: Scan for USB devices
echo "Step 3: Scanning for USB storage devices..."
echo ""

# Arrays to store device information
declare -a DEVICES
declare -a LABELS
declare -a UUIDS
declare -a SIZES
declare -a FSTYPES

# Scan for removable storage devices
INDEX=0
while IFS= read -r line; do
    # Parse lsblk output: NAME SIZE TYPE FSTYPE UUID LABEL MOUNTPOINT
    NAME=$(echo "$line" | awk '{print $1}')
    SIZE=$(echo "$line" | awk '{print $2}')
    TYPE=$(echo "$line" | awk '{print $3}')
    FSTYPE=$(echo "$line" | awk '{print $4}')
    UUID=$(echo "$line" | awk '{print $5}')
    LABEL=$(echo "$line" | awk '{print $6}')

    # Only include partitions with filesystem and UUID
    if [[ "$TYPE" == "part" ]] && [[ ! -z "$UUID" ]] && [[ ! -z "$FSTYPE" ]]; then
        # Check if it's a removable device
        DISK_NAME=$(echo "$NAME" | sed 's/[0-9]*$//')
        if [ -f "/sys/block/$DISK_NAME/removable" ]; then
            REMOVABLE=$(cat "/sys/block/$DISK_NAME/removable")
            if [ "$REMOVABLE" == "1" ]; then
                DEVICES[$INDEX]="/dev/$NAME"
                SIZES[$INDEX]="$SIZE"
                FSTYPES[$INDEX]="$FSTYPE"
                UUIDS[$INDEX]="$UUID"
                LABELS[$INDEX]="${LABEL:--}"
                INDEX=$((INDEX + 1))
            fi
        fi
    fi
done < <(lsblk -nlo NAME,SIZE,TYPE,FSTYPE,UUID,LABEL | grep -E "sd[a-z][0-9]|mmcblk[0-9]p[0-9]")

# Display found devices
if [ $INDEX -eq 0 ]; then
    echo -e "${RED}✗ No removable USB drives found${NC}"
    echo ""
    echo "Please ensure USB drives are connected and try again"
    exit 1
fi

echo -e "${GREEN}✓ Found $INDEX removable USB drive(s)${NC}"
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo -e "${CYAN}Available USB Drives:${NC}"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

for i in "${!DEVICES[@]}"; do
    echo ""
    echo -e "${BLUE}[$i]${NC} ${DEVICES[$i]}"
    echo "    Label:  ${LABELS[$i]}"
    echo "    Size:   ${SIZES[$i]}"
    echo "    Type:   ${FSTYPES[$i]}"
    echo "    UUID:   ${UUIDS[$i]}"
done

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""

# Step 4: Select drives to mount
echo "Step 4: Select drives to auto-mount"
echo ""

declare -a SELECTED_DEVICES
declare -a SELECTED_UUIDS
declare -a SELECTED_MOUNT_POINTS

read -p "How many drives do you want to auto-mount? (1-$INDEX): " NUM_DRIVES

if ! [[ "$NUM_DRIVES" =~ ^[0-9]+$ ]] || [ "$NUM_DRIVES" -lt 1 ] || [ "$NUM_DRIVES" -gt $INDEX ]; then
    echo -e "${RED}✗ Invalid number${NC}"
    exit 1
fi

echo ""

for ((n=0; n<$NUM_DRIVES; n++)); do
    echo -e "${CYAN}Drive $(($n + 1)) of $NUM_DRIVES:${NC}"

    # Select device
    while true; do
        read -p "  Select device number [0-$((INDEX-1))]: " DEV_NUM
        if [[ "$DEV_NUM" =~ ^[0-9]+$ ]] && [ "$DEV_NUM" -ge 0 ] && [ "$DEV_NUM" -lt $INDEX ]; then
            break
        fi
        echo -e "  ${RED}Invalid selection, try again${NC}"
    done

    # Select mount point
    read -p "  Mount point name (e.g., usb0, usb1, backup): " MOUNT_NAME
    MOUNT_POINT="/media/$MOUNT_NAME"

    SELECTED_DEVICES[$n]="${DEVICES[$DEV_NUM]}"
    SELECTED_UUIDS[$n]="${UUIDS[$DEV_NUM]}"
    SELECTED_MOUNT_POINTS[$n]="$MOUNT_POINT"

    echo -e "  ${GREEN}✓ Will mount ${DEVICES[$DEV_NUM]} to $MOUNT_POINT${NC}"
    echo ""
done

# Step 5: Confirm
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo -e "${CYAN}Configuration Summary:${NC}"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
for i in "${!SELECTED_DEVICES[@]}"; do
    echo ""
    echo "Drive $(($i + 1)):"
    echo "  Device:      ${SELECTED_DEVICES[$i]}"
    echo "  UUID:        ${SELECTED_UUIDS[$i]}"
    echo "  Mount Point: ${SELECTED_MOUNT_POINTS[$i]}"
done
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""

read -p "Proceed with setup? (y/n): " -n 1 -r
echo ""

if ! [[ $REPLY =~ ^[Yy]$ ]]; then
    echo "Setup cancelled"
    exit 0
fi

echo ""

# Step 6: Disable fstab entries
echo "Step 6: Backing up and disabling conflicting fstab entries..."
cp /etc/fstab /etc/fstab.backup.$(date +%Y%m%d_%H%M%S)

# Comment out any existing USB mount entries
for mount_point in "${SELECTED_MOUNT_POINTS[@]}"; do
    mount_name=$(basename "$mount_point")
    sed -i "/$mount_name/s/^[^#]/#DISABLED_BY_SYSTEMD# /" /etc/fstab
done

echo -e "${GREEN}✓ Fstab backed up and updated${NC}"
echo ""

# Step 7: Create systemd units
echo "Step 7: Creating systemd mount units..."
echo ""

create_mount_unit() {
    local uuid=$1
    local mount_point=$2
    local fstype=$3
    local unit_name=$(systemd-escape -p --suffix=mount "$mount_point")

    # Determine mount options based on filesystem type
    local mount_opts="defaults,nofail,uid=$REAL_UID,gid=$REAL_GID"
    if [[ "$fstype" == "vfat" ]] || [[ "$fstype" == "exfat" ]]; then
        mount_opts+=",umask=0022"
    fi

    cat > "/etc/systemd/system/$unit_name" <<EOF
[Unit]
Description=USB Drive Mount ($mount_point)
After=local-fs.target

[Mount]
What=UUID=$uuid
Where=$mount_point
Type=$fstype
Options=$mount_opts

[Install]
WantedBy=multi-user.target
EOF

    # Create automount unit
    local automount_unit="${unit_name/.mount/.automount}"
    cat > "/etc/systemd/system/$automount_unit" <<EOF
[Unit]
Description=USB Drive Auto-mount Trigger ($mount_point)

[Automount]
Where=$mount_point
TimeoutIdleSec=5

[Install]
WantedBy=multi-user.target
EOF

    echo -e "${GREEN}✓ Created units for $mount_point${NC}"
}

# Create mount points and units for each selected drive
for i in "${!SELECTED_DEVICES[@]}"; do
    mount_point="${SELECTED_MOUNT_POINTS[$i]}"
    uuid="${SELECTED_UUIDS[$i]}"

    # Detect filesystem type
    dev_index=-1
    for j in "${!DEVICES[@]}"; do
        if [ "${DEVICES[$j]}" == "${SELECTED_DEVICES[$i]}" ]; then
            dev_index=$j
            break
        fi
    done

    fstype="vfat"
    if [ $dev_index -ge 0 ]; then
        fstype="${FSTYPES[$dev_index]}"
    fi

    # Create mount point
    mkdir -p "$mount_point"
    chown $REAL_UID:$REAL_GID "$mount_point"

    # Create systemd units
    create_mount_unit "$uuid" "$mount_point" "$fstype"
done

echo ""

# Step 8: Reload systemd and enable units
echo "Step 8: Enabling auto-mount..."
systemctl daemon-reload

for mount_point in "${SELECTED_MOUNT_POINTS[@]}"; do
    mount_name=$(basename "$mount_point")
    unit_name=$(systemd-escape -p --suffix=automount "$mount_point")

    systemctl enable "$unit_name"
    systemctl start "$unit_name"

    echo -e "${GREEN}✓ Enabled $mount_point auto-mount${NC}"
done

echo ""
echo "=========================================="
echo "Setup Complete!"
echo "=========================================="
echo ""
echo "USB drives will now auto-mount on access"
echo ""
echo "Test your mounts:"
for mount_point in "${SELECTED_MOUNT_POINTS[@]}"; do
    echo "  ls $mount_point"
done
echo ""
echo "Check status:"
for mount_point in "${SELECTED_MOUNT_POINTS[@]}"; do
    mount_name=$(basename "$mount_point")
    unit_name=$(systemd-escape -p --suffix=automount "$mount_point")
    echo "  systemctl status $unit_name"
done
echo ""
echo "View mounted drives:"
echo "  mount | grep media"
echo ""