#!/bin/bash
#
# Complete USB Auto-mount Setup
# Cleans up existing mounts and sets up systemd automount
#

set -e

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

echo "=========================================="
echo "Complete 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_final.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-usb0.automount 2>/dev/null || true
systemctl stop media-usb1.automount 2>/dev/null || true
systemctl stop media-usb0.mount 2>/dev/null || true
systemctl stop media-usb1.mount 2>/dev/null || true
systemctl disable media-usb0.automount 2>/dev/null || true
systemctl disable media-usb1.automount 2>/dev/null || true
echo -e "${GREEN}✓ Stopped existing units${NC}"
echo ""

# Step 2: Unmount USB drives
echo "Step 2: Unmounting USB drives..."
umount /media/usb0 2>/dev/null || true
umount /media/usb1 2>/dev/null || true
echo -e "${GREEN}✓ Unmounted drives${NC}"
echo ""

# Step 3: Extract UUIDs from fstab (including commented lines)
echo "Step 3: Reading UUIDs from /etc/fstab..."
USB0_UUID=$(grep "usb0" /etc/fstab 2>/dev/null | sed 's/^#[^ ]* //' | grep -oP 'UUID=\K[^\s]+' || echo "")
USB1_UUID=$(grep "usb1" /etc/fstab 2>/dev/null | sed 's/^#[^ ]* //' | grep -oP 'UUID=\K[^\s]+' || echo "")

if [ -z "$USB0_UUID" ] && [ -z "$USB1_UUID" ]; then
    echo -e "${YELLOW}⚠ No UUIDs found in fstab${NC}"
    echo ""
    echo "Please enter UUIDs manually:"
    read -p "USB0 UUID (e.g., D20B-AEC0) [Enter to skip]: " USB0_UUID
    read -p "USB1 UUID (e.g., E212-4316) [Enter to skip]: " USB1_UUID
else
    echo -e "${GREEN}✓ Found USB0 UUID: $USB0_UUID${NC}"
    echo -e "${GREEN}✓ Found USB1 UUID: $USB1_UUID${NC}"
fi
echo ""

# Step 4: Comment out fstab entries to prevent conflicts
echo "Step 4: Disabling fstab entries (creating backup)..."
cp /etc/fstab /etc/fstab.backup.$(date +%Y%m%d_%H%M%S)
sed -i '/UUID=.*usb[01]/s/^/#DISABLED_BY_SYSTEMD# /' /etc/fstab
echo -e "${GREEN}✓ Fstab entries commented out${NC}"
echo ""

# Step 5: Create mount points
echo "Step 5: 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 ""

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

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

    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=vfat
Options=defaults,nofail,uid=$REAL_UID,gid=$REAL_GID,umask=0022

[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}"
}

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 "${RED}✗ No UUIDs provided, cannot create units${NC}"
    exit 1
fi

echo ""

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

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 on access"
echo ""
echo "Test with:"
echo "  ls /media/usb0"
echo "  ls /media/usb1"
echo ""
echo "Check status:"
echo "  systemctl status media-usb0.automount"
echo "  systemctl status media-usb1.automount"
echo ""
echo "View mounted drives:"
echo "  mount | grep usb"
echo ""
echo "Note: Original fstab backed up with timestamp"
echo "Fstab entries commented out to prevent conflicts"
echo ""