#!/bin/sh

# A UI for detecting and selecting all displays. Probes xrandr for connected
# displays and lets user select one to use. User may also select "manual
# selection" which opens arandr.
#
# Options:
#   -r    Enable resolution selection for each display

# Parse options
select_resolution=false
while getopts "r" opt; do
  case $opt in
    r) select_resolution=true ;;
    *) echo "Usage: $0 [-r]" >&2; exit 1 ;;
  esac
done
shift $((OPTIND - 1))

# Function to select resolution for a display
get_resolution() {
  display="$1"
  if [ "$select_resolution" = true ]; then
    resolutions=$(xrandr --query | sed -n "/^$display/,/^[^ ]/p" | grep -v "^$display" | grep -v "^[^ ]" | awk '{print $1}' | grep -E "^[0-9]+x[0-9]+$")
    chosen_res=$(echo "$resolutions" | dmenu -i -p "Resolution for $display:")
    [ -n "$chosen_res" ] && echo "--mode $chosen_res" || echo "--auto"
  else
    # No EDID-preferred mode on some displays makes --auto fall back to a tiny
    # mode (e.g. 640x480). Default to the highest available mode instead.
    best=$(xrandr --query | sed -n "/^$display/,/^[^ ]/p" |
      grep -v "^$display" | grep -E "^[[:space:]]+[0-9]+x[0-9]+" |
      awk '{print $1}' | sort -t x -k1,1n -k2,2n | tail -n1)
    [ -n "$best" ] && echo "--mode $best" || echo "--auto"
  fi
}

# Map a friendly direction to the correct xrandr relative-position flag.
# Note: left/right use the "-of" suffix, but above/below do NOT.
dir_flag() {
  case "$1" in
    left) echo "--left-of" ;;
    right) echo "--right-of" ;;
    above) echo "--above" ;;
    below) echo "--below" ;;
  esac
}

# Return the opposite direction (used to place a third display).
opposite_dir() {
  case "$1" in
    left) echo "right" ;;
    right) echo "left" ;;
    above) echo "below" ;;
    below) echo "above" ;;
  esac
}

# Extract "WIDTH HEIGHT" for a display given its resolution option
# ("--mode WxH" or "--auto"). Falls back to the display's highest mode.
res_wh() {
  display="$1"
  opt="$2"
  case "$opt" in
  *--mode*) echo "$opt" | sed 's/.*--mode //; s/x/ /' ;;
  *)
    xrandr --query | sed -n "/^$display/,/^[^ ]/p" |
      grep -v "^$display" | grep -E "^[[:space:]]+[0-9]+x[0-9]+" |
      awk '{print $1}' | sort -t x -k1,1n -k2,2n | tail -n1 | sed 's/x/ /'
    ;;
  esac
}

twoscreen() { # If multi-monitor is selected and there are two screens.
  mirror=$(printf "no\\nyes" | dmenu -i -p "Mirror displays?")
  [ -z "$mirror" ] && exit
  # Mirror displays using native resolution of external display and a scaled
  # version for the internal display
  if [ "$mirror" = "yes" ]; then
    external=$(echo "$screens" | dmenu -i -p "Optimize resolution for:")
    internal=$(echo "$screens" | grep -v "$external")

    res_external=$(xrandr --query | sed -n "/^$external/,/\+/p" |
      tail -n 1 | awk '{print $1}')
    res_internal=$(xrandr --query | sed -n "/^$internal/,/\+/p" |
      tail -n 1 | awk '{print $1}')

    res_ext_x=$(echo "$res_external" | sed 's/x.*//')
    res_ext_y=$(echo "$res_external" | sed 's/.*x//')
    res_int_x=$(echo "$res_internal" | sed 's/x.*//')
    res_int_y=$(echo "$res_internal" | sed 's/.*x//')

    scale_x=$(echo "$res_ext_x / $res_int_x" | bc -l)
    scale_y=$(echo "$res_ext_y / $res_int_y" | bc -l)

    xrandr --output "$external" --auto --scale 1.0x1.0 \
      --output "$internal" --auto --same-as "$external" \
      --scale "$scale_x"x"$scale_y"
  else
    primary=$(echo "$screens" | dmenu -i -p "Select primary display:")
    [ -z "$primary" ] && exit
    secondary=$(echo "$screens" | grep -v ^"$primary"$)
    direction=$(printf "left\\nright\\nabove\\nbelow" | dmenu -i -p "What side of $primary should $secondary be on?")
    [ -z "$direction" ] && exit
    primary_res=$(get_resolution "$primary")
    secondary_res=$(get_resolution "$secondary")
    case "$direction" in
    left | right)
      xrandr --output "$primary" --primary $primary_res --scale 1.0x1.0 \
        --output "$secondary" $(dir_flag "$direction") "$primary" $secondary_res --scale 1.0x1.0
      ;;
    above | below)
      # --above/--below left-align by default; compute positions to right-align.
      set -- $(res_wh "$primary" "$primary_res"); pw=$1; ph=$2
      set -- $(res_wh "$secondary" "$secondary_res"); sw=$1; sh=$2
      maxw=$pw; [ "$sw" -gt "$maxw" ] && maxw=$sw
      px=$((maxw - pw)); sx=$((maxw - sw))
      if [ "$direction" = "above" ]; then py=$sh; sy=0; else py=0; sy=$ph; fi
      xrandr --output "$primary" --primary --mode "${pw}x${ph}" --pos "${px}x${py}" --scale 1.0x1.0 \
        --output "$secondary" --mode "${sw}x${sh}" --pos "${sx}x${sy}" --scale 1.0x1.0
      ;;
    esac
  fi
}

morescreen() { # If multi-monitor is selected and there are more than two screens.
  primary=$(echo "$screens" | dmenu -i -p "Select primary display:")
  [ -z "$primary" ] && exit
  secondary=$(echo "$screens" | grep -v ^"$primary"$ | dmenu -i -p "Select secondary display:")
  direction=$(printf "left\\nright\\nabove\\nbelow" | dmenu -i -p "What side of $primary should $secondary be on?")
  [ -z "$direction" ] && exit
  tertiary=$(echo "$screens" | grep -v ^"$primary"$ | grep -v ^"$secondary"$ | dmenu -i -p "Select third display:")
  primary_res=$(get_resolution "$primary")
  secondary_res=$(get_resolution "$secondary")
  tertiary_res=$(get_resolution "$tertiary")
  case "$direction" in
  left | right)
    xrandr --output "$primary" --primary $primary_res \
      --output "$secondary" $(dir_flag "$direction") "$primary" $secondary_res \
      --output "$tertiary" $(dir_flag "$(opposite_dir "$direction")") "$primary" $tertiary_res
    ;;
  above | below)
    # Stack vertically (secondary on the chosen side, tertiary opposite) and
    # right-align all three to the widest display.
    set -- $(res_wh "$primary" "$primary_res"); pw=$1; ph=$2
    set -- $(res_wh "$secondary" "$secondary_res"); sw=$1; sh=$2
    set -- $(res_wh "$tertiary" "$tertiary_res"); tw=$1; th=$2
    maxw=$pw; [ "$sw" -gt "$maxw" ] && maxw=$sw; [ "$tw" -gt "$maxw" ] && maxw=$tw
    px=$((maxw - pw)); sx=$((maxw - sw)); tx=$((maxw - tw))
    if [ "$direction" = "above" ]; then
      sy=0; py=$sh; ty=$((sh + ph)) # secondary top, primary middle, tertiary bottom
    else
      ty=0; py=$th; sy=$((th + ph)) # tertiary top, primary middle, secondary bottom
    fi
    xrandr --output "$primary" --primary --mode "${pw}x${ph}" --pos "${px}x${py}" \
      --output "$secondary" --mode "${sw}x${sh}" --pos "${sx}x${sy}" \
      --output "$tertiary" --mode "${tw}x${th}" --pos "${tx}x${ty}"
    ;;
  esac
}

multimon() { # Multi-monitor handler.
  case "$(echo "$screens" | wc -l)" in
  2) twoscreen ;;
  *) morescreen ;;
  esac
}

onescreen() { # If only one output available or chosen.
  res_opt=$(get_resolution "$1")
  xrandr --output "$1" --primary $res_opt --scale 1.0x1.0 $(echo "$allposs" | grep -v "\b$1" | awk '{print "--output", $1, "--off"}' | paste -sd ' ' -)
}

postrun() { # Stuff to run to clean up.
  setbg     # Fix background if screen size/arangement has changed.
  {
    killall dunst
    setsid -f dunst
  } >/dev/null 2>&1 # Restart dunst to ensure proper location on screen
}

# Get all possible displays
allposs=$(xrandr -q | grep "connected")

# Get all connected screens.
screens=$(echo "$allposs" | awk '/ connected/ {print $1}')

# If there's only one screen
[ "$(echo "$screens" | wc -l)" -lt 2 ] &&
  {
    onescreen "$screens"
    postrun
    notify-send "💻 Only one screen detected." "Using it in its optimal settings..."
    exit
  }

# Get user choice including multi-monitor and manual selection:
chosen=$(printf "%s\\nmulti-monitor\\nmanual selection" "$screens" | dmenu -i -p "Select display arangement:") &&
  case "$chosen" in
  "manual selection")
    arandr
    exit
    ;;
  "multi-monitor") multimon ;;
  *) onescreen "$chosen" ;;
  esac

postrun