#!/bin/bash

# USB Auto-Mount Script for Raspberry Pi Video Player
# This script automatically mounts USB devices when they are inserted

# Configuration
MOUNT_BASE="/media/usb"
LOG_FILE="/var/log/usb_automount.log"

# Load environment variables from .env file if it exists
if [[ -f ".env" ]]; then
    source .env
fi

# Set defaults if not defined in .env
USER="${USER:-pi}"

# Function to log messages
log_message() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" >> "$LOG_FILE"
}

# Function to create mount point
create_mount_point() {
    local device_name="$1"
    local mount_point="$MOUNT_BASE/$device_name"
    
    if [ ! -d "$mount_point" ]; then
        mkdir -p "$mount_point"
        chown "$USER:$USER" "$mount_point"
        log_message "Created mount point: $mount_point"
    fi
}

# Function to mount USB device
mount_usb_device() {
    local device_path="$1"
    local device_name=$(basename "$device_path")
    local mount_point="$MOUNT_BASE/$device_name"
    
    # Check if device is already mounted
    if mountpoint -q "$mount_point" 2>/dev/null; then
        log_message "Device $device_name is already mounted at $mount_point"
        return 0
    fi
    
    # Create mount point
    create_mount_point "$device_name"
    
    # Get device label (if available)
    local label=$(blkid -s LABEL -o value "$device_path" 2>/dev/null)
    if [ -n "$label" ]; then
        # Use label as mount point name if available
        mount_point="$MOUNT_BASE/$label"
        create_mount_point "$label"
    fi
    
    # Mount the device
    if mount -o uid=1000,gid=1000,umask=0002 "$device_path" "$mount_point" 2>/dev/null; then
        log_message "Successfully mounted $device_name at $mount_point"
        
        # Create symlink for easy access
        ln -sf "$mount_point" "/home/$USER/Desktop/USB-$device_name" 2>/dev/null || true
        
        # Notify user (if desktop environment is running)
        if [ -n "$DISPLAY" ]; then
            notify-send "USB Device Mounted" "Device $device_name mounted at $mount_point" 2>/dev/null || true
        fi
        
        return 0
    else
        log_message "Failed to mount $device_name at $mount_point"
        return 1
    fi
}

# Function to unmount USB device
unmount_usb_device() {
    local device_path="$1"
    local device_name=$(basename "$device_path")
    
    # Find mount point
    local mount_point=$(mount | grep "$device_path" | awk '{print $3}')
    
    if [ -n "$mount_point" ]; then
        if umount "$mount_point" 2>/dev/null; then
            log_message "Successfully unmounted $device_name from $mount_point"
            
            # Remove symlink
            rm -f "/home/$USER/Desktop/USB-$device_name" 2>/dev/null || true
            
            # Remove mount point if empty
            if [ -d "$mount_point" ] && [ -z "$(ls -A "$mount_point" 2>/dev/null)" ]; then
                rmdir "$mount_point" 2>/dev/null || true
            fi
            
            # Notify user
            if [ -n "$DISPLAY" ]; then
                notify-send "USB Device Unmounted" "Device $device_name has been unmounted" 2>/dev/null || true
            fi
            
            return 0
        else
            log_message "Failed to unmount $device_name from $mount_point"
            return 1
        fi
    else
        log_message "Device $device_name is not mounted"
        return 1
    fi
}

# Function to mount all existing USB devices
mount_existing_devices() {
    log_message "Checking for existing USB devices..."
    
    # Find all USB block devices
    for device in /dev/sd*[0-9]; do
        if [ -b "$device" ]; then
            # Check if it's a USB device
            if readlink -f "$device" | grep -q "usb"; then
                mount_usb_device "$device"
            fi
        fi
    done
}

# Function to scan for video files on mounted USB devices
scan_usb_videos() {
    log_message "Scanning USB devices for video files..."
    
    local video_extensions="mp4|avi|mkv|mov|wmv|flv|webm|m4v"
    local video_count=0
    
    for mount_point in "$MOUNT_BASE"/*; do
        if [ -d "$mount_point" ] && mountpoint -q "$mount_point" 2>/dev/null; then
            local videos=$(find "$mount_point" -type f -iregex ".*\.\($video_extensions\)" 2>/dev/null | wc -l)
            if [ "$videos" -gt 0 ]; then
                log_message "Found $videos video files on $(basename "$mount_point")"
                video_count=$((video_count + videos))
            fi
        fi
    done
    
    log_message "Total video files found on USB devices: $video_count"
}

# Main function
main() {
    case "$1" in
        "add")
            if [ -n "$2" ]; then
                mount_usb_device "$2"
            else
                echo "Usage: $0 add <device_path>"
                exit 1
            fi
            ;;
        "remove")
            if [ -n "$2" ]; then
                unmount_usb_device "$2"
            else
                echo "Usage: $0 remove <device_path>"
                exit 1
            fi
            ;;
        "scan")
            scan_usb_videos
            ;;
        "mount-all")
            mount_existing_devices
            ;;
        *)
            echo "Usage: $0 {add|remove|scan|mount-all} [device_path]"
            echo "  add <device>    - Mount a USB device"
            echo "  remove <device> - Unmount a USB device"
            echo "  scan           - Scan for video files on mounted USB devices"
            echo "  mount-all      - Mount all existing USB devices"
            exit 1
            ;;
    esac
}

# Create necessary directories
mkdir -p "$MOUNT_BASE"
mkdir -p "$(dirname "$LOG_FILE")"

# Run main function with all arguments
main "$@"