#!/usr/bin/env bash
## Grimblast: a helper for screenshots within hyprland
## Requirements:
##  - `grim`: screenshot utility for wayland
##  - `slurp`: to select an area
##  - `hyprctl`: to read properties of current window
##  - `wl-copy`: clipboard utility
##  - `jaq`: json utility to parse hyprctl output
##  - `notify-send`: to show notifications
## Those are needed to be installed, if unsure, run `grimblast check`
##
## See `man 1 grimblast` or `grimblast usage` for further details.

## Author: Misterio (https://github.com/misterio77)

## This tool is based on grimshot, with swaymsg commands replaced by their
## hyprctl equivalents.
## https://github.com/swaywm/sway/blob/master/contrib/grimshot

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}}"
  echo "$HOME/Screenshots"
}

NOTIFY=no
CURSOR=

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

  case $key in
    -n|--notify)
      NOTIFY=yes
      shift # past argument
      ;;
    -c|--cursor)
      CURSOR=yes
      shift # past argument
      ;;
    *)    # unknown option
      break # done with parsing --flags
      ;;
  esac
done

ACTION=${1:-usage}
SUBJECT=${2:-screen}
FILE=${3:-$(getTargetDirectory)/$(date +%Y-%m-%d-%H-%M-%S).png}

if [ "$ACTION" != "save" ] && [ "$ACTION" != "copy" ] && [ "$ACTION" != "copysave" ] && [ "$ACTION" != "check" ]; then
  echo "Usage:"
  echo "  grimblast [--notify] [--cursor] (copy|save|copysave) [active|screen|output|area] [FILE|-]"
  echo "  grimblast check"
  echo "  grimblast 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 "  copysave: Combine the previous 2 options."
  echo "  check: Verify if required tools are installed and exit."
  echo "  usage: Show this message and exit."
  echo ""
  echo "Targets:"
  echo "  active: Currently active window."
  echo "  screen: All visible outputs."
  echo "  output: Currently active output."
  echo "  area: Manually select a region or window."
  exit
fi

notify() {
  if [ $FILE = "" ]; then
    notify-send -t 3000 -a grimblast "$@"
  else
    notify-send -t 3000 -i $FILE -a grimblast "$@"
  fi
}
notifyOk() {
  [ "$NOTIFY" = "no" ] && return

  TITLE=${2:-"Screenshot"}
  MESSAGE=${1:-"OK"}
  REST=${@: 3}
  notify "$TITLE" "$MESSAGE" $REST
}
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 hyprctl
  check wl-copy
  check jaq
  check notify-send
  exit
elif [ "$SUBJECT" = "active" ] ; then
  FOCUSED=$(hyprctl activewindow -j)
  GEOM=$(echo "$FOCUSED" | jaq -r '"\(.at[0]),\(.at[1]) \(.size[0])x\(.size[1])"')
  APP_ID=$(echo "$FOCUSED" | jaq -r '.class')
  WHAT="$APP_ID window"
elif [ "$SUBJECT" = "screen" ] ; then
  GEOM=""
  WHAT="Screen"
elif [ "$SUBJECT" = "output" ] ; then
  GEOM=""
  OUTPUT=$(hyprctl monitors -j | jaq -r '.[] | select(.focused == true)' | jaq -r '.name')
  WHAT="$OUTPUT"
elif [ "$SUBJECT" = "area" ] ; then
  WORKSPACES="$(hyprctl monitors -j | jaq -r 'map(.activeWorkspace.id)')"
  WINDOWS="$(hyprctl clients -j | jaq -r --argjson workspaces "$WORKSPACES" 'map(select([.workspace.id] | inside($workspaces)))' )"
  GEOM=$(echo "$WINDOWS" | jaq -r '.[] | "\(.at[0]),\(.at[1]) \(.size[0])x\(.size[1])"' | slurp)
  # Check if user exited slurp without selecting the area
  if [ -z "$GEOM" ]; then
   exit 1
  fi
  WHAT="Area"
elif [ "$SUBJECT" = "window" ] ; then
  die "Subject 'window' is now included in 'area'"
else
  die "Unknown subject to take a screen shot from" "$SUBJECT"
fi

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