#!/bin/bash
#
# USB Drive Setup Script
# Mounts USB drives with proper user permissions for wedding phone backups
#

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 Drive Setup for Wedding Phone"
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.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 USB mounts for user: $REAL_USER (UID:$REAL_UID, GID:$REAL_GID)"
echo ""

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

# Detect USB devices
echo "Detecting USB devices..."
lsblk -o NAME,SIZE,TYPE,MOUNTPOINT,FSTYPE | grep -E "disk|part"
echo ""

# Function to mount a USB drive
mount_usb() {
    local device=$1
    local mount_point=$2

    if [ ! -b "$device" ]; then
        echo -e "${YELLOW}⚠ Device $device not found, skipping${NC}"
        return
    fi

    # Unmount if already mounted
    if mountpoint -q "$mount_point"; then
        echo "Unmounting existing mount at $mount_point..."
        umount "$mount_point" 2>/dev/null || true
    fi

    echo "Mounting $device to $mount_point..."
    mount -o uid=$REAL_UID,gid=$REAL_GID "$device" "$mount_point"

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✓ Mounted $device${NC}"

        # Test write permissions
        if sudo -u $REAL_USER touch "$mount_point/.test" 2>/dev/null; then
            sudo -u $REAL_USER rm "$mount_point/.test"
            echo -e "${GREEN}✓ Write test successful${NC}"
        else
            echo -e "${RED}✗ Write test failed${NC}"
        fi

        # Show available space
        df -h "$mount_point" | tail -n 1 | awk '{print "  Free space: " $4}'
    else
        echo -e "${RED}✗ Failed to mount $device${NC}"
    fi

    echo ""
}

# Prompt for devices
echo "Enter device paths or press Enter to skip"
echo ""

read -p "USB 0 device (e.g., /dev/sda1) [Enter to skip]: " USB0_DEVICE
if [ ! -z "$USB0_DEVICE" ]; then
    mount_usb "$USB0_DEVICE" "/media/usb0"
fi

read -p "USB 1 device (e.g., /dev/sdb1) [Enter to skip]: " USB1_DEVICE
if [ ! -z "$USB1_DEVICE" ]; then
    mount_usb "$USB1_DEVICE" "/media/usb1"
fi

# Summary
echo "=========================================="
echo "Mount Summary"
echo "=========================================="
ls -la /media/usb0 /media/usb1 2>/dev/null | grep -E "^d|^total" || echo "No mounts found"
echo ""

# Offer to add to fstab
read -p "Add these mounts to /etc/fstab for automatic mounting? (y/n) " -n 1 -r
echo ""

if [[ $REPLY =~ ^[Yy]$ ]]; then
    echo ""
    echo "Add the following lines to /etc/fstab:"
    echo ""

    if [ ! -z "$USB0_DEVICE" ]; then
        UUID0=$(blkid -s UUID -o value "$USB0_DEVICE")
        if [ ! -z "$UUID0" ]; then
            echo "UUID=$UUID0  /media/usb0  vfat  defaults,nofail,uid=$REAL_UID,gid=$REAL_GID  0  0"
        fi
    fi

    if [ ! -z "$USB1_DEVICE" ]; then
        UUID1=$(blkid -s UUID -o value "$USB1_DEVICE")
        if [ ! -z "$UUID1" ]; then
            echo "UUID=$UUID1  /media/usb1  vfat  defaults,nofail,uid=$REAL_UID,gid=$REAL_GID  0  0"
        fi
    fi

    echo ""
    echo "To add automatically, run:"
    echo "  sudo nano /etc/fstab"
fi

echo ""
echo "=========================================="
echo "Setup Complete!"
echo "=========================================="
echo ""
echo "USB drives are now accessible by user: $REAL_USER"
echo "Wedding Phone can now backup to these drives"
echo ""