#!/bin/sh

## This is a modification of Grimshot: a helper for screenshots within sway,
## for use with River compositor. It supports only "screen" (full screen) and
## "area" (selected area) commands

## Requirements:
##  - `grim`: screenshot utility for wayland
##  - `slurp`: to select an area
##  - `wl-copy`: clipboard utility
##  - `notify-send`: to show notifications
## Those are needed to be installed, if unsure, run `river-grimshot check`

getTargetDirectory() {
  test -f "${XDG_CONFIG_HOME:-$HOME/.config}/user-dirs.dirs" && \
    . "${XDG_CONFIG_HOME:-$HOME/.config}/user-dirs.dirs"

  echo "${XDG_SCREENSHOTS_DIR:-${XDG_PICTURES_DIR:-$HOME}}"
}

NOTIFY=no
CURSOR=
WAIT=no

while [ $# -gt 0 ]; do
  key="$1"

  case $key in
    -n|--notify)
      NOTIFY=yes
      shift # past argument
      ;;
    -c|--cursor)
      CURSOR=yes
      shift # past argument
      ;;
    -w|--wait)
      shift
      WAIT="$1"
      if echo "$WAIT" | grep "[^0-9]" -q; then
        echo "invalid value for wait '$WAIT'" >&2
        exit 3
      fi
      shift
      ;;
    *)    # unknown option
      break # done with parsing --flags
      ;;
  esac
done

ACTION=${1:-usage}
SUBJECT=${2:-screen}
FILE=${3:-$(getTargetDirectory)/$(date -Ins).png}

if [ "$ACTION" != "save" ] && [ "$ACTION" != "copy" ] && [ "$ACTION" != "check" ]; then
  echo "Usage:"
  echo "  river-grimshot [--notify] [--cursor] [--wait N] (copy|save) [screen|area] [FILE|-]"
  echo "  river-grimshot check"
  echo "  river-grimshot usage"
  echo ""
  echo "Commands:"
  echo "  copy: Copy the screenshot data into the clipboard."
  echo "  save: Save the screenshot to a regular file or '-' to pipe to STDOUT."
  echo "  check: Verify if required tools are installed and exit."
  echo "  usage: Show this message and exit."
  echo ""
  echo "Targets:"
  echo "  screen: All visible outputs."
  echo "  area: Manually select a region."
  exit
fi

notify() {
  notify-send -t 3000 -a river-grimshot "$@"
}
notifyOk() {
  [ "$NOTIFY" = "no" ] && return

  TITLE=${2:-"Screenshot"}
  MESSAGE=${1:-"OK"}
  notify "$TITLE" "$MESSAGE"
}
notifyError() {
  if [ $NOTIFY = "yes" ]; then
    TITLE=${2:-"Screenshot"}
    MESSAGE=${1:-"Error taking screenshot with grim"}
    notify -u critical "$TITLE" "$MESSAGE"
  else
    echo "$1"
  fi
}

die() {
  MSG=${1:-Bye}
  notifyError "Error: $MSG"
  exit 2
}

check() {
  COMMAND=$1
  if command -v "$COMMAND" > /dev/null 2>&1; then
    RESULT="OK"
  else
    RESULT="NOT FOUND"
  fi
  echo "   $COMMAND: $RESULT"
}

takeScreenshot() {
  FILE=$1
  GEOM=$2
  OUTPUT=$3
  if [ -n "$OUTPUT" ]; then
    grim ${CURSOR:+-c} -o "$OUTPUT" "$FILE" || die "Unable to invoke grim"
  elif [ -z "$GEOM" ]; then
    grim ${CURSOR:+-c} "$FILE" || die "Unable to invoke grim"
  else
    grim ${CURSOR:+-c} -g "$GEOM" "$FILE" || die "Unable to invoke grim"
  fi
}

if [ "$ACTION" = "check" ] ; then
  echo "Checking if required tools are installed. If something is missing, install it to your system and make it available in PATH..."
  check grim
  check slurp
  check wl-copy
  check notify-send
  exit
elif [ "$SUBJECT" = "area" ] ; then
  GEOM=$(slurp -d)
  # Check if user exited slurp without selecting the area
  if [ -z "$GEOM" ]; then
    exit 1
  fi
  WHAT="Area"
elif [ "$SUBJECT" = "screen" ] ; then
  GEOM=""
  WHAT="Screen"
else
  die "Unknown subject to take a screen shot from" "$SUBJECT"
fi

if [ "$WAIT" != "no" ]; then
  sleep "$WAIT"
fi

if [ "$ACTION" = "copy" ] ; then
  takeScreenshot - "$GEOM" "$OUTPUT" | wl-copy --type image/png || die "Clipboard error"
  notifyOk "$WHAT copied to buffer"
else
  if takeScreenshot "$FILE" "$GEOM" "$OUTPUT"; then
    TITLE="Screenshot of $SUBJECT"
    MESSAGE=$(basename "$FILE")
    notifyOk "$MESSAGE" "$TITLE"
    echo "$FILE"
  else
    notifyError "Error taking screenshot with grim"
  fi
fi