#!/bin/bash

# Define lock file
LOCK_FILE="/tmp/usb_bind.lock"

# Function to check if vmid is a number between 100 and 200
is_numeric_vmid() {
    local vmid=$1
    [[ "$vmid" =~ ^[0-9]+$ ]] && (( vmid >= 100 && vmid < 200 ))
}

# Validate format xxxx:xxxx
is_valid_pidvid() {
    [[ "$1" =~ ^[0-9a-fA-F]{4}:[0-9a-fA-F]{4}$ ]]
}

is_valid_label() {
    local label="$1"
    [[ "$(lsusb | grep -i "$label")" ]]
}

# Check USB device speed
check_usb_speed() {
    local vendor_id=$1
    local product_id=$2
    lsusb -d "$vendor_id:$product_id" -v 2>/dev/null | grep -q "Bus" || { echo "Device $vendor_id:$product_id not connected"; return 0; }
    usb_version=$(lsusb -d "$vendor_id:$product_id" -v 2>/dev/null | grep "bcdUSB" | awk '{print $2}')
    speed_supported=$(lsusb -d "$vendor_id:$product_id" -v 2>/dev/null | grep -i "wSpeedsSupported" | awk '{print $2}')
    [[ "$usb_version" == "3.00" && "$speed_supported" =~ "000e" ]] && return 2
    return 1
}

# ----------- Show help if no arguments provided -------------
if [[ $# -lt 2 ]]; then
    echo "Usage: $0 <VM ID or VM Name> <PID:VID or USB Label>"
    echo
    echo "Arguments:"
    echo "  VM ID or Name   - The VM ID or Name (numeric ID or '::next' for next available VM)"
    echo "    Example: 101 or 'myVMName' or '::next'"
    echo "  PID:VID         - The USB device identifier in format 'PID:VID' (e.g. 1234:5678)"
    echo "    Example: 1234:5678"
    echo "  USB Label       - A string (case-insensitive) representing the USB device label (e.g. 'razer', 'keyboard')"
    echo "    Example: 'razer mouse' or 'gaming keyboard'"
    echo
    echo "Options:"
    echo "  --skip          - Skip if the device is already binded to a running vm"
    echo "    Example: --skip"
    echo "  --alsobind PID:VID - Also bind the USB device to another VM, can specify another PID:VID or USB Label"
    echo "    Example: --alsobind 1234:5678 or --alsobind 'razer mouse'"
    echo "  --vmlist=vmid1,vmid2,... - Only used with '::next', a comma-separated list of VM IDs (numeric only). Only valid VM IDs that are running will be considered."
    echo "    Example: --vmlist=101,102,103"
    echo "  --nolock        - Skip lock checking for recursive calls (useful for recursive script calls)"
    echo "    Example: --nolock"
    echo
    echo "Examples:"
    echo "  $0 101 1234:5678             # Bind USB device 1234:5678 to VM 101"
    echo "  $0 'myVMName' 'razer mouse'  # Bind the USB device with label 'razer mouse' to 'myVMName'"
    echo "  $0 '::next' 1234:5678        # Bind USB device 1234:5678 to the next available VM"
    echo "  $0 101 'gaming keyboard' --alsobind 1234:5678 --alsobind 6789:1254  # Bind 'gaming keyboard',1234:5678 and 6789:1254  to VM 101"
    echo
    exit 1
fi
# ----------- ARGUMENTS PARSING -------------
vmid_or_name=$1
device="$2"
shift 2

if ! is_valid_pidvid "$device" && ! is_valid_label "$device"; then
    echo "Invalid device format or label: $device"
    exit 1
fi

if is_valid_pidvid "$device"; then
    pid="${device%%:*}"
    vid="${device##*:}"
    # proceed as before
elif is_valid_label "$device"; then
    device_info=$(lsusb | grep -i "$device" | head -n 1 | awk '{print $6}')
    pid="${device_info%%:*}"
    vid="${device_info##*:}"
else
    echo "Device not found"
    exit 1
fi

skip_delete=false
alsobinds=()
custom_vm_list=()
nolock=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        --skip) skip_delete=true ;;
        --alsobind)
            shift
            if ! is_valid_pidvid "$1" && ! is_valid_label "$1"; then
                echo "Invalid alsobind device : $1"
                exit 1
            fi
            alsobinds+=("$1")
            ;;
        --vmlist=*)
            IFS=',' read -ra raw_vm_list <<< "${1#--vmlist=}"
            for vm in "${raw_vm_list[@]}"; do
                if [[ "$vm" =~ ^[0-9]+$ ]]; then
                    # Get the VM status and ensure it's running
                    vm_status=$(qm list | awk -v id="$vm" '$1 == id {print $3}')
                    [[ "$vm_status" == "running" ]] && custom_vm_list+=("$vm")
                else
                    echo "Invalid VMID: $vm. VMIDs must be numeric."
                    exit 1
                fi
            done
            ;;
        --nolock) nolock=true ;;
    esac
    shift
done

if [[ "$nolock" == false ]]; then

    # Check if the lock file exists
    if [ -f "$LOCK_FILE" ]; then
        echo "Another instance of this script is already running. Exiting."
        exit 1
    fi

    # Create lock file to indicate that the script is running
    touch "$LOCK_FILE"

    # Ensure the lock file is removed when the script exits
    trap "rm -f $LOCK_FILE" EXIT


fi




# Resolve VMID
if is_numeric_vmid "$vmid_or_name"; then
    vmid=$vmid_or_name
elif [[ "$vmid_or_name" == "0" || "${vmid_or_name,,}" == "::none" || "${vmid_or_name,,}" == "::next" ]]; then
    vmid=""
else
    vmid=$(qm list | awk -F' ' -v name="$vmid_or_name" 'tolower($2) == tolower(name) {print $1}')
    [[ -z "$vmid" ]] && exit 1
fi

# Check if VM is running
vm_status=$(qm list | awk -v vmid="$vmid" '$1 == vmid {print $3}')
[[ -n "$vmid" && "$vm_status" != "running" ]] && exit 1

# USB Speed check
check_usb_speed "$pid" "$vid"
usb_status=$?
[[ $usb_status -eq 0 ]] && exit 1
is_usb3=""
[[ $usb_status -eq 2 ]] && is_usb3=",usb3=yes"

# Analyse VMs
vms=($(qm list | grep 'running' | awk '{print $1, $2}'))
if [[ ${#vms[@]} -eq 0 ]]; then
    echo "No running VMs."
    exit 1
fi
vm_usb_used=false
fix_problem_slot=""
need_fix_problem=false
free_slot=""

if [[ "${vmid_or_name,,}" == "::next" ]]; then
    if [[ ${#custom_vm_list[@]} -gt 0 ]]; then
        # Custom list defined
        found_index=-1
        for ((i=0; i<${#custom_vm_list[@]}; i++)); do
            vmid_candidate="${custom_vm_list[$i]}"
            vm_config=$(qm config "$vmid_candidate" 2>/dev/null)
            if echo "$vm_config" | grep -q "host=$pid:$vid"; then
                found_index=$i
                break
            fi
        done

        if [[ $found_index -ge 0 ]]; then
            next_index=$(( (found_index + 1) % ${#custom_vm_list[@]} ))
        else
            next_index=0
        fi
        vmid="${custom_vm_list[$next_index]}"
    else
        # Default list: all running VMs
        found_index=-1
        for ((i=0; i<${#vms[@]}; i+=2)); do
            vmid_running="${vms[$i]}"
            vm_config=$(qm config "$vmid_running")
            if echo "$vm_config" | grep -q "host=$pid:$vid"; then
                found_index=$i
                break
            fi
        done

        if [[ $found_index -ge 0 ]]; then
            next_index=$((found_index + 2))
            if [[ $next_index -ge ${#vms[@]} ]]; then
                next_index=0
            fi
            vmid="${vms[$next_index]}"
        else
            vmid="${vms[0]}"
        fi
    fi
fi

for ((i=0; i<${#vms[@]}; i+=2)); do
    vmid_running="${vms[$i]}"
    vmname="${vms[$i+1]}"
    vm_config=$(qm config "$vmid_running")

    if echo "$vm_config" | grep -q "host=$pid:$vid"; then
        slot=$(echo "$vm_config" | grep "host=$pid:$vid" | awk '{print $1}' | sed 's/:$//')
        if [[ "$vmid_running" == "$vmid" ]]; then
            vm_usb_used=true
            fix_problem_slot="$slot"
        else
            [[ "$skip_delete" == false ]] && qm set "$vmid_running" -delete "$slot"
            need_fix_problem=true
        fi
    fi

    if [[ "$vmid_running" == "$vmid" ]]; then
        for z in {0..13}; do
            slot="usb$z"
            if ! echo "$vm_config" | grep -q "$slot"; then
                free_slot="$slot"
                break
            fi
        done
    fi
done

if [[ "$need_fix_problem" == true && "$vm_usb_used" == true && "$skip_delete" == false ]]; then
    qm set "$vmid" -delete "$fix_problem_slot"
    vm_usb_used=false
fi

if [[ -n "$free_slot" && "$vm_usb_used" == false && -n "$vmid" ]]; then
    if [[ "$need_fix_problem" == true && "$skip_delete" == true ]]; then
        echo "skip because already attributed"
        exit 1
    else
        qm set "$vmid" --$free_slot host=$pid:$vid$is_usb3
    fi
fi

also_bind_target="$vmid"
if [[ -z "$also_bind_target" ]]; then
    also_bind_target="0"
fi

# ----------- RECURSIVE CALLS FOR --alsobind ------------
for extra in "${alsobinds[@]}"; do
    if is_valid_pidvid "$extra"; then
        extra_pid="${extra%%:*}"
        extra_vid="${extra##*:}"
    elif is_valid_label "$extra"; then
        extra_info=$(lsusb | grep -i "$extra" | awk '{print $6}')
        extra_pid="${extra_info%%:*}"
        extra_vid="${extra_info##*:}"
    fi
    sleep 3
    echo "Also binding: $also_bind_target $extra -> $extra_pid:$extra_vid"
    "$0" "$also_bind_target" "$extra" --nolock
done