#!/bin/bash
# Author: Philipp Schmitt
# Version: 2.7.3
# Dependencies: xrandr awk sed dzen2

XRDR_CONF_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/xrdr"
HOOK_FILE='/etc/conf.d/xrdr/xrdr.hook'
USER_HOOK_FILE="$XRDR_CONF_DIR/xrdr.hook"

if [[ -r $USER_HOOK_FILE ]]; then
    HOOK=$USER_HOOK_FILE
elif [[ -r $HOOK_FILE ]]; then
    HOOK=$HOOK_FILE
fi

usage() {
    echo "Usage: $1 [extend|copy|primary|secondary|tertiary|layout]"
    echo "            [rotate ANGLE DISPLAY]"
    echo "extend: Extend displays (3 displays max)"
    echo "copy: Copy output (2 displays max)"
    echo "primary: Disable all except primary display"
    echo "secondary: Disable all except second display"
    echo "tertiary: Disable all except third display"
    echo "rotate: Rotate a display (by default +90)"
    echo "layout [FILE]: Apply layout from file (or read from stdin)"
    echo "print: Print current layout to stdout"
}

# TODO Support rotation + primary display layouts ("1 2*r 3")
parse_config() {
    [[ -n $VERBOSE ]] && echo "Applying layout from $1"

    local valid_config=$(sed '/^#\|^$/d;s/|\|;\|,/ /g;s/#.*//' <<< "$1")
    local config
    local rows=$(wc -l <<< "$valid_config")
    local row=0 row_content desired_screens=0

    [[ -n $VERBOSE ]] && echo -e "Layout to be applied: $valid_config"

    for row in $(seq $rows); do
        row_content=$(sed -n ${row}p <<< "$valid_config")
        eval "local config_${row}=( $row_content )"
        eval "let desired_screens=desired_screens+\${#config_$row[@]}"
        eval "local config[$row]=\${config_${row}[@]}"
    done

    [[ $desired_screens -gt ${#SCREENS[@]} ]] && {
        echo "There aren't that many screens attached"
        exit 1
    }

    local config_row config_row_index=0
    local _screen screen screen_index previous_screen
    local xrandr_cmd="xrandr --output"

    for config_row in $(seq ${#config[@]}); do
        eval "row_content=(\${config_$config_row[*]})"
        let config_row_index++
        for _screen in "${row_content[@]}"; do
            screen_index=$(sed 's/[^0-9]*//g' <<< "$_screen")
            screen_attributes=$(sed 's/[0-9]*//g' <<< "$_screen")
            screen=${SCREENS[$(( $screen_index - 1 ))]}
            xrandr_cmd="$xrandr_cmd $screen --auto"
            case $screen_attributes in
                r)  xrandr_cmd="$xrandr_cmd --rotate right"    ;;
                l)  xrandr_cmd="$xrandr_cmd --rotate left"     ;;
                i)  xrandr_cmd="$xrandr_cmd --rotate inverted" ;;
                n)  xrandr_cmd="$xrandr_cmd --rotate normal"   ;;
                x)  xrandr_cmd="$xrandr_cmd --off"             ;;
                \*) xrandr_cmd="$xrandr_cmd --primary"         ;;
            esac
            if [[ $_screen == ${row_content[0]} ]]; then
                if [[ $config_row_index -ne 1 ]]; then
                    xrandr_cmd="$xrandr_cmd --below $first_screen"
                fi
                first_screen=$screen
            else
                xrandr_cmd="$xrandr_cmd --right-of $previous_screen"
            fi
            if [[ $_screen != ${row_content[${#row_content[@]} - 1]} \
               || $config_row_index -lt ${#config[@]} ]]; then
                xrandr_cmd="$xrandr_cmd --output"
                previous_screen=$screen
            fi
        done
    done

    [[ -n $VERBOSE ]] && echo $xrandr_cmd
    [[ -z $DRY_RUN ]] && {
        disconnect
        $xrandr_cmd
    } || echo $xrandr_cmd
    exit 0
}

_get_screen_at_pos() {
    local x_pos=$1 y_pos=$2
    for screen in ${!SCREENS[@]}; do
        read _ _ x y <<< $(_res $screen)
       [[ $x == $x_pos ]] && [[ $y == $y_pos ]] && {
            echo $screen
            return
        }
    done
}

_get_config_size() {
    local rows cols
    local xpos ypos
    for screen in ${!SCREENS[@]}; do
        read _ _ xpos ypos <<< $(_res $screen)
        rows="$rows\n$xpos"
        cols="$cols\n$ypos"
    done

    cols=$(( $(echo -e "$cols" | sort -u | wc -l) - 1))
    rows=$(( $(echo -e "$rows" | sort -u | wc -l) - 1))
    echo "$rows $cols"
}

_get_y_positions() {
    local screen ypos rows
    for screen in ${!SCREENS[@]}; do
        read _ _ _ ypos <<< $(_res $screen)
        [[ -n $rows ]] && rows="$rows\n$ypos" || rows="$ypos"
    done
    echo -e "$rows" | sort -u | tr '\n' ' '
}

_get_screens_per_row() {
    local screens=0
    local rows=( $(_get_y_positions) )
    local ypos

    [[ "${rows[@]}" =~ $1 ]] || { echo "$1 not in ${rows[*]}" 2>&1; exit 1; }
    for screen in ${!SCREENS[@]}; do
        read _ _ _ ypos <<< $(_res $screen)
        [[ $ypos -eq ${rows[$1]} ]] && let screens++
    done
    echo $screens
}

_get_x_positions_in_row() {
    local screen xpos ypos rows
    for screen in ${!SCREENS[@]}; do
        read _ _ xpos ypos <<< $(_res $screen)
        [[ $ypos -eq $1 ]] && {
            [[ -n $rows ]] && rows="$rows\n$xpos" || rows="$xpos"
        }
    done
    echo -e "$rows" | sort -u | tr '\n' ' '
}

print_layout() {
    # TODO Improve performance
    # TODO Detect primary display
    [[ -n $VERBOSE ]] && echo "Printing layout to stdout"

    local maxcols screens_per_row c r ypos_values xpos_values
    local screen o layout tmp
    read _ maxcols <<< $(_get_config_size)
    ypos_values=( $(_get_y_positions) )
    for (( c = 0; c < $maxcols; c++ )); do
        screens_in_current_row=$(_get_screens_per_row $c)
        xpos_values=( $(_get_x_positions_in_row ${ypos_values[$c]}) )
        for (( r = 0; r < $screens_in_current_row; r++ )); do
            read xres yres xpos ypos <<< $(_res $screen)
            screen=$(_get_screen_at_pos ${xpos_values[$r]} ${ypos_values[$c]})
            o=$(_orient $screen)
            tmp="$(( $screen + 1))$o"
            [[ -n $layout ]] && {
                [[ $r -ne 0 ]] && {
                    layout="${layout} | $tmp"
                } || {
                    layout="${layout}${tmp}"
                }
            } || {
                layout="$tmp"
            }
        done
        [[ $c != $(( maxcols - 1 )) ]] && layout="$layout\n"
    done

    echo -e $layout
}

_check_rotation() {
    [[ $1 == "left" ]] \
    || \
    [[ $1 == "right" ]] \
    || \
    [[ $1 == "normal" ]] \
    || \
    [[ $1 == "inverted" ]]
}

_guess_screens() {
    PRIMARY_SCREEN=${SCREENS[0]}
    SECONDARY_SCREEN=${SCREENS[1]}
    TERTIARY_SCREEN=${SCREENS[2]}
}

count() {
    echo ${#SCREENS[@]}
    exit 0
}

identify() {
    local x_res y_res screen screen_num
    for screen in ${!SCREENS[@]}; do
        screen_num=$(( $screen + 1 ))
        read x_res y_res _ _ <<< $(_res $screen)
        echo "Screen $screen_num (${SCREENS[$screen]}) - ${x_res}x${y_res}" \
            | dzen2 -p 2 -xs $screen_num -bg red -w $x_res -h $y_res
    done
    exit 0
}

# This function returns ${XRES} ${YRES} ${XPOS} ${YPOS}
_res() {
    if [[ $1 -ge ${#SCREENS[@]} ]]; then
        exit 1
    fi
    xrandr | grep ${SCREENS[$1]} \
           | sed -e 's/.* \([0-9]\+x[0-9]\++[0-9]\++[0-9]\+\) .*/\1/' \
                 -e 's/x/ /g;s/+/ /g'
}

resolution() {
    local x_res y_res
    read x_res y_res _ _ <<< $(_res $1)
    echo $x_res - $y_res
}

width() {
    _res $1 | awk '{ print $1 }'
}

height() {
    _res $1 | awk '{ print $2 }'
}

_orient() {
    local o=$(xrandr | grep ${SCREENS[$1]} | awk '{ print $4 }')
    case $o in
        # normal)   echo n ;;
        left)     echo l ;;
        right)    echo r ;;
        inverted) echo i ;;
        *)        echo   ;;
    esac
}

is_vertical() {
    # TODO awk '{ print $4 }' will only work if $1 != primary screen
    local vert=$(xrandr | grep ${SCREENS[$1]} | awk '{ print $4}')
    exit $([[ $vert == "right" || $vert == "left" ]])
}

copy() {
    disconnect
    xrandr --output $SECONDARY_SCREEN --auto \
           --output $PRIMARY_SCREEN --same-as $SECONDARY_SCREEN --primary
}

extend() {
    disconnect
    case ${#SCREENS[@]} in
        2)
            xrandr --output $PRIMARY_SCREEN --primary --auto \
                   --output $SECONDARY_SCREEN --auto --right-of $PRIMARY_SCREEN
            ;;
        3)
            xrandr --output $PRIMARY_SCREEN --primary --auto \
                   --output $SECONDARY_SCREEN --auto --left-of $PRIMARY_SCREEN \
                   --output $TERTIARY_SCREEN --auto --right-of $PRIMARY_SCREEN
            ;;
    esac
}

primary_only() {
    disconnect
    case ${#SCREENS[@]} in
        3)
            xrandr --output $TERTIARY_SCREEN --off
            ;&
        2)
            xrandr --output $SECONDARY_SCREEN --off
            ;;
    esac
    xrandr --output $PRIMARY_SCREEN --auto --primary
}

secondary_only() {
    [[ -n "$SECONDARY_SCREEN" ]] || exit 1
    disconnect
    xrandr --output $PRIMARY_SCREEN --off \
           --output $TERTIARY_SCREEN --off \
           --output $SECONDARY_SCREEN --auto --primary
}

tertiary_only() {
    [[ -n "$TERTIARY_SCREEN" ]] || exit 1
    disconnect
    xrandr --output $PRIMARY_SCREEN --off \
           --output $SECONDARY_SCREEN --off \
           --output $TERTIARY_SCREEN --auto --primary
}

auto_rotate() {
    local o=$(xrandr | grep ${SCREENS[$1]} | awk '{ print $4 }')
    echo $o
    local new_o
    case $o in
        left)     new_o=inverted ;;
        right)    new_o=normal   ;;
        inverted) new_o=right    ;;
        *)        new_o=left     ;;
    esac
    rotate ${SCREENS[$1]} $new_o
}

rotate() {
    xrandr --output $1 --rotate $2
    _hook $2
}

disconnect() {
    xrandr --auto
}

toggle() {
    [[ ${#SCREENS[@]} -gt 1 ]] && primary_only || extend
}

auto() {
    [[ ${#SCREENS[@]} -gt 1 ]] && extend || primary_only
}

_hook() {
    [[ -r $HOOK ]] && . $HOOK "$*" || return 0
}

while getopts 'd:hnsvx:' opt ; do
    case $opt in
        d)
            # FIXME is the export necessary?
            export DISPLAY="$OPTARG"
            ;;
        h)
            usage $(basename $0)
            exit 0
            ;;
        n)
            DRY_RUN=1
            ;;
        s)
            NOSORT=1
            ;;
        v)
            VERBOSE=1
            ;;
        x)
            # FIXME is the export necessary?
            export XAUTHORITY="$OPTARG"
            ;;
        \?)
            echo "Invalid option: -$OPTARG" >&2
            exit 1
            ;;
        :)
            echo "Option -$OPTARG requires an argument." >&2
            exit 1
            ;;
    esac
done

shift $(($OPTIND - 1))

[[ -z "$DISPLAY" ]] && export DISPLAY=:0
[[ -z "$XAUTHORITY" ]] && export XAUTHORITY=$HOME/.Xauthority

[[ -z "$NOSORT" ]] && {
    SCREENS=($(xrandr | sed -ne 's/\(.\+\) connected.*/\1/p' | sort -u | tr '\n' ' '))
} || {
    SCREENS=($(xrandr | sed -ne 's/\(.\+\) connected.*/\1/p' | tr '\n' ' '))
}
_guess_screens

case "$1" in
    auto|a)
        action=auto
        ;;
    count|c)
        count
        ;;
    primary|p)
        primary_only
        ;;
    extend|e)
        extend
        ;;
    copy|cp)
        copy
        ;;
    secondary|s)
        secondary_only
        ;;
    tertiary|ter)
        tertiary_only
        ;;
    toggle|t)
        toggle
        ;;
    rotate|r)
        action=rot
        ;;
    identify|i)
        identify
        ;;
    res|resolution)
        action=r
        ;;
    height|h)
        action=h
        ;;
    width|w)
        action=w
        ;;
    vertical|v)
        action=v
        ;;
    help|h)
        usage $(basename $0)
        exit 0
        ;;
    l|layout)
        action=l
        shift
        ;;
    print)
        print_layout
        exit 0
        ;;
    *)
        action=auto
        ;;
esac

[[ -n $action ]] && {
    [[ -n $2 && $2 = *[[:digit:]]* ]] && {
        [[ $2 -gt ${#SCREENS[@]} || $2 -lt 0 ]] && {
            echo "Invalid index"
            exit 1
        } || {
            d=$(( $2 - 1 ))
        }
        shift
    }
    # Make sure $d is set
    [[ -z "$d" ]] && d=${SCREENS[0]}
    case $action in
        auto) auto ;;
        rot)
        [[ -z "$2" ]] && auto_rotate $d || {
            _check_rotation $2 && {
                rotate $d $2
            } || exit 1
        }
        ;;
        r) resolution $d  ;;
        w) width $d       ;;
        h) height $d      ;;
        v) is_vertical $d ;;
        l)
            [[ $# -eq 0 ]] && {
                echo "Reading layout from stdin"
                read config
            } || {
                [[ -r $1 ]] || { echo "File not readable"; exit 1; }
                config=$(cat $1)
            }
            parse_config "$config"
        ;;
    esac
} || _hook

# vim: set ft=sh et ts=4 sw=4 cc=80 : #
