#!/bin/bash
#
# USB Auto-mount Setup using systemd
# This creates systemd mount units that auto-mount when USB drives are inserted
#

set -e

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

echo "=========================================="
echo "USB Auto-mount Setup (systemd)"
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_automount.sh"
    exit 1
fi

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

echo "Setting up auto-mount for user: $REAL_USER (UID:$REAL_UID, GID:$REAL_GID)"
echo ""

# Create mount points
echo "Creating mount points..."
mkdir -p /media/usb0 /media/usb1
chown $REAL_UID:$REAL_GID /media/usb0 /media/usb1
echo -e "${GREEN}✓ Mount points created${NC}"
echo ""

# Detect USB devices and get UUIDs
echo "Detecting USB devices..."
lsblk -f | grep -E "sd[a-z][0-9]|vfat|exfat|ntfs"
echo ""

# Function to create systemd mount unit
create_mount_unit() {
    local uuid=$1
    local mount_point=$2
    local unit_name=$(systemd-escape -p --suffix=mount "$mount_point")

    echo "Creating systemd mount unit: $unit_name"

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

[Mount]
What=UUID=$uuid
Where=$mount_point
Type=vfat
Options=defaults,nofail,uid=$REAL_UID,gid=$REAL_GID,umask=0022

[Install]
WantedBy=multi-user.target
EOF

    # Create corresponding 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 $unit_name and $automount_unit${NC}"
}

# Get UUIDs from existing fstab or user input
USB0_UUID=""
USB1_UUID=""

# Try to extract UUIDs from fstab
if grep -q "UUID=.*usb0" /etc/fstab; then
    USB0_UUID=$(grep "usb0" /etc/fstab | grep -oP 'UUID=\K[^\s]+')
    echo "Found USB0 UUID in fstab: $USB0_UUID"
fi

if grep -q "UUID=.*usb1" /etc/fstab; then
    USB1_UUID=$(grep "usb1" /etc/fstab | grep -oP 'UUID=\K[^\s]+')
    echo "Found USB1 UUID in fstab: $USB1_UUID"
fi

echo ""
echo "If UUIDs were not detected, please enter them manually:"
echo ""

if [ -z "$USB0_UUID" ]; then
    read -p "USB0 UUID (e.g., D20B-AEC0) [Enter to skip]: " USB0_UUID
fi

if [ -z "$USB1_UUID" ]; then
    read -p "USB1 UUID (e.g., E212-4316) [Enter to skip]: " USB1_UUID
fi

# Create systemd units
echo ""
CREATED=0

if [ ! -z "$USB0_UUID" ]; then
    create_mount_unit "$USB0_UUID" "/media/usb0"
    CREATED=1
fi

if [ ! -z "$USB1_UUID" ]; then
    create_mount_unit "$USB1_UUID" "/media/usb1"
    CREATED=1
fi

if [ $CREATED -eq 0 ]; then
    echo -e "${YELLOW}⚠ No UUIDs provided, no units created${NC}"
    exit 0
fi

# Reload systemd
echo ""
echo "Reloading systemd daemon..."
systemctl daemon-reload
echo -e "${GREEN}✓ Systemd reloaded${NC}"

# Enable and start automount units
echo ""
echo "Enabling auto-mount units..."

if [ ! -z "$USB0_UUID" ]; then
    systemctl enable media-usb0.automount
    systemctl start media-usb0.automount
    echo -e "${GREEN}✓ Enabled /media/usb0 auto-mount${NC}"
fi

if [ ! -z "$USB1_UUID" ]; then
    systemctl enable media-usb1.automount
    systemctl start media-usb1.automount
    echo -e "${GREEN}✓ Enabled /media/usb1 auto-mount${NC}"
fi

echo ""
echo "=========================================="
echo "Setup Complete!"
echo "=========================================="
echo ""
echo "USB drives will now auto-mount when accessed or plugged in"
echo ""
echo "To test:"
echo "  ls /media/usb0"
echo "  ls /media/usb1"
echo ""
echo "To check status:"
echo "  systemctl status media-usb0.automount"
echo "  systemctl status media-usb1.automount"
echo ""
echo "To disable (if needed):"
echo "  sudo systemctl disable media-usb0.automount"
echo "  sudo systemctl disable media-usb1.automount"
echo ""