#!/bin/bash

# This script is used to apply monitor profiles based on the connected monitors and USB devices.

set -euo pipefail

# Get the number of monitors connected to the system
get_monitor_count() {
    xrandr --listmonitors | head -n1 | cut -d' ' -f2
}

# Get the list of connected and disconnected monitors.
# Includes information about the primary monitor
# E.g:
# eDP-1,connected,primary
# HDMI-1,disconnected
# DP-1,disconnected
_get_monitors() {
    xrandr | grep -Eo '^.* ((dis|)connected)( primary|) ' | sed 's/ /,/g' | sed 's/,$//g'
}

# Function that gets all monitors and takes an optional amount of arguments to filter the list.
# Filters are applied as an OR operation.
# E.g:
# get_monitors connected primary
# get_monitors connected primary eDP-1
get_monitors() {
    local monitors
    monitors=$(_get_monitors)

    # If no filters are set, return all monitors
    if [ "$#" -eq 0 ]; then
        echo "$monitors"
        return 0
    fi

    for monitor in $monitors; do
        for filter in "$@"; do
            for part in ${monitor//,/ }; do
                if [ "$part" = "$filter" ]; then
                    echo "$monitor"
                    continue 3 # Continue to the next monitor
                fi
            done
        done
    done
}

# monitor_exist checks if the specified monitor is connected to the system
# $1: The monitor name to search for
monitor_exist() {
    local monitor=${1:?Monitor not set}
    local connected_monitors
    connected_monitors=$(get_monitors connected)

    for connected_monitor in $connected_monitors; do
        if [ "$(cut -d',' -f1 <<<"$connected_monitor")" = "$monitor" ]; then
            return 0
        fi
    done

    return 1
}

# monitors_exist checks if the specified monitors are connected to the system
# $1..$n: The monitor names to search for
monitors_exist() {
    local monitor matches=0

    for monitor in "$@"; do
        if monitor_exist "$monitor"; then
            echo "$monitor"
            matches=$((matches + 1))
        fi
    done

    if [ "$matches" -ne "$#" ]; then
        echo "Amount of devices does not match the amount of arguments" 1>&2
        return 1
    fi

    return 0
}

# Disable all monitors that are not in the list of connected monitors
disable_disconnected_monitors() {
    local connected_monitors
    connected_monitors=$(get_monitors disconnected)

    echo "Disabling disconnected monitors"

    for monitor in $connected_monitors; do
        echo "Disabling monitor $(cut -d',' -f1 <<<"$monitor")"
    done
}

# Get the list of USB devices using lsusb, outputs for example "06cb:00f0 Synaptics, Inc."
get_usb_devices() {
    lsusb | grep -Eo ' ID .*' | cut -d' ' -f3-
}

# Get the list of USB device names using lsusb, outputs for example "Synaptics, Inc."
get_usb_device_names() {
    get_usb_devices | cut -d' ' -f2-
}

# Get the list of USB devices that match the name
get_usb_devices_by_name() {
    local name=${1:?Name not set}
    get_usb_device_names | while read -r device; do
        if [ "$device" = "$name" ]; then
            echo "$device"
        fi
    done
}

# Get the list of USB devices that match the pattern
get_usb_devices_by_pattern() {
    local pattern=${1:?Pattern not set}
    get_usb_device_names | grep -E "$pattern"
}

# Returns the list of USB devices that match the criteria
# $1: Required, either "all" or "any"
# $2: Mode, either "exact" or "pattern"
# $3..$n: The name or pattern to search for
usb_devices_exist() {
    local required=${1:?Required not set}
    local mode=${2:?Mode not set}
    shift 2
    local devices matches=0

    # Check lenght of arguments
    if [ "$#" -eq 0 ]; then
        echo "No patterns set" 1>&2
        return 1
    fi

    # Check if required and mode args are valid
    if ! [[ "$required" =~ ^(all|any)$ ]]; then
        echo "Invalid 'required' argument: $required" 1>&2
        return 1
    elif ! [[ "$mode" =~ ^(exact|pattern)$ ]]; then
        echo "Invalid 'mode' argument: $mode" 1>&2
        return 1
    fi

    # Process each pattern
    for pattern in "$@"; do
        if [ -z "$pattern" ]; then
            echo "Got empty pattern" 1>&2
            return 1
        fi

        if [ "$mode" = "exact" ]; then
            devices="$(get_usb_devices_by_name "$pattern")"
        elif [ "$mode" = "pattern" ]; then
            devices="$(get_usb_devices_by_pattern "$pattern")"
        fi

        # Check if at least one device was found
        if [ -n "$devices" ]; then
            matches=$((matches + 1))
            echo "$devices"
        fi
    done

    # Check if all patterns had a match
    if [ "$required" = "all" ]; then
        if [ "$matches" -eq "$#" ]; then
            return 0
        else
            echo "Amount of devices does not match the amount of patterns" 1>&2
            return 1
        fi
    # Check if at least one pattern had a match
    elif [ "$required" = "any" ]; then
        if [ "$matches" -ge "0" ]; then
            return 0
        else
            echo "Amount of devices does not match the amount of patterns" 1>&2
            return 1
        fi
    fi
}

# Get the list of empty i3 workspaces
get_empty_i3_workspaces() {
    i3-msg -t get_tree | jq -r '
        recurse(.nodes[]?) |                # Recursively descend into "nodes" arrays
        select(.type == "workspace") |      # Select elements where "type" is "workspace"
        select(.name | test("^[^_].*")) |   # Select elements where "name" does not start with "_"
        select(.nodes | length == 0) |      # Further select those where "nodes" array is empty
        .name                               # Output the "name" of these workspaces
    '
}

i3_workspace_is_empty() {
    local workspace_name=${1:?Workspace name not set}
    local empty_workspace
    empty_workspace=$(get_empty_i3_workspaces)

    for workspace in $empty_workspace; do
        if [ "$workspace" = "$workspace_name" ]; then
            return 0
        fi
    done

    return 1
}

# Get a comma separated list of i3 workspace names and the output they are on
get_i3_workspaces() {
    i3-msg -t get_workspaces | jq -r '.[] | .name + "," + .output'
}

get_i3_workspace_prop() {
    local workspace_name=${1:?Workspace name not set}
    local prop=${2:?Property not set}
    i3-msg -t get_workspaces | jq -r ".[] | select(.name == \"$workspace_name\") | .$prop"
}

i3_workspace_exists() {
    local workspace_name=${1:?Workspace name not set}
    local workspaces name

    if [[ -n "$(get_i3_workspace_prop "$workspace_name" name)" ]]; then
        return 0
    fi

    return 1
}

# This function selects a workspace by name and moves it to the desired monitor.
# If the workspace name is an integer, it will allow incrementing the workspace name by the given number
move_i3_workspace() {
    local workspace_name=${1:?Workspace name not set}
    local monitor_name=${2:?Monitor name not set}
    local increment=${3:-0}
    local new_workspace_name="$workspace_name"

    if ! i3_workspace_exists "$workspace_name"; then
        echo "Workspace $workspace_name not found" 1>&2
        return 1
    elif ! monitor_exist "$monitor_name"; then
        echo "Monitor $monitor_name not found" 1>&2
        return 1
    fi

    # Ensure the new output is not the same as the current output
    if [[ "$(get_i3_workspace_prop "$workspace_name" output)" == "$monitor_name" ]]; then
        echo "Workspace $workspace_name is already on monitor $monitor_name" 1>&2
        return 1
    fi

    # Increment the workspace name if it is an integer
    if [[ "$workspace_name" =~ ^[0-9]+$ ]]; then
        new_workspace_name=$((workspace_name + increment))
    elif [ "$increment" -ne 0 ]; then
        echo "Workspace name is not an integer. Unable to increment/decrement." 1>&2
        return 1
    fi

    # Check if the new workspace name already exists and is not the same as the current workspace name
    if i3_workspace_exists "$new_workspace_name"; then
        if [[ "$new_workspace_name" != "$workspace_name" ]]; then
            echo "Workspace $new_workspace_name already exists" 1>&2
            return 1
        fi
    fi

    echo "Moving workspace $workspace_name to monitor $monitor_name as $new_workspace_name"

    i3-msg "workspace $workspace_name" 1>/dev/null # Select the workspace
    i3-msg "rename workspace $workspace_name to $new_workspace_name" 1>/dev/null
    i3-msg "move workspace to output $monitor_name" 1>/dev/null
}

# This function moves all workspaces that are not on the given monitor to the given monitor
move_nonexistent_i3_workspaces() {
    local monitor_name=${1:?Monitor name not set}
    local increment=${2:-0}
    local workspaces
    workspaces=$(get_i3_workspaces)

    echo "Moving workspaces to monitor $monitor_name"

    for workspace in $workspaces; do
        local workspace_name
        local workspace_output
        workspace_name=$(echo "$workspace" | cut -d',' -f1)
        workspace_output=$(echo "$workspace" | cut -d',' -f2)

        if i3_workspace_is_empty "$workspace_name"; then
            echo "Skipping empty workspace $workspace_name"
            continue
        fi

        if ! monitor_exist "$workspace_output"; then
            move_i3_workspace "$workspace_name" "$monitor_name" "$increment"
        fi
    done
}

main() {
    local primary_monitor monitor_count monitors=() profile_found=false exists
    primary_monitor=$(get_monitors primary | cut -d',' -f1)
    monitor_count=$(get_monitors connected | wc -l)
    echo "Monitors connected: $monitor_count"

    if [ "$monitor_count" -eq 1 ]; then
        echo "Applying laptop-only profile"
        xrandr --auto
        profile_found=true
    elif usb_devices_exist all exact "HP, Inc HP USB-C Universal Dock" >/dev/null 2>&1; then
        monitors=(
            "eDP-1,DVI-I-1-1,DVI-I-2-2"
            "eDP-1,DVI-I-2-2,DVI-I-3-3"
        )

        for monitor in "${monitors[@]}"; do
            IFS=',' read -r -a monitor_array <<<"$monitor"
            if monitors_exist "${monitor_array[@]}" >/dev/null; then
                echo "Applying work profile for ${monitor_array[*]}"
                xrandr \
                    --output "${monitor_array[0]}" --mode 1920x1200 --rotate normal --pos 1200x1474 --primary \
                    --output "${monitor_array[1]}" --mode 1920x1200 --rotate left --pos 0x0 \
                    --output "${monitor_array[2]}" --mode 1920x1200 --rotate normal --pos 1200x274
                profile_found=true
                break
            fi
        done
    elif usb_devices_exist all exact "Lenovo ThinkPad USB-C Dock Audio" >/dev/null 2>&1; then
        monitor_array=(eDP-1 DP-1-0.3 DP-1-0.1)
        if monitors_exist "${monitor_array[@]}" >/dev/null; then
            echo "Applying home profile for ${monitor_array[*]}"
            xrandr \
                --output "${monitor_array[0]}" --mode 2560x1440 --rotate normal --pos 0x480 --primary \
                --output "${monitor_array[1]}" --mode 3440x1440 --rotate normal --pos 2560x480 \
                --output "${monitor_array[2]}" --mode 1920x1200 --rotate right --pos 6000x0
            profile_found=true
        fi

        monitor_array=(DP-3-3 eDP-1 DP-3-1)
        if monitors_exist "${monitor_array[@]}" >/dev/null; then
            xrandr \
                --output "${monitor_array[0]}" --mode 3440x1440 --pos 0x480 --rotate normal \
                --output "${monitor_array[1]}" --mode 1920x1200 --pos 3440x720 --rotate normal --primary \
                --output "${monitor_array[2]}" --mode 1920x1200 --pos 5360x0 --rotate right
            profile_found=true
        fi
    fi

    if ! "$profile_found"; then
        echo "No profile found"
        return 1
    fi

    disable_disconnected_monitors
    move_nonexistent_i3_workspaces "$primary_monitor" 10
    i3-msg restart >/dev/null
}

# Run if invoked directly
if [ "$0" = "${BASH_SOURCE[0]}" ]; then
    main
fi