#!/bin/bash

# Take a screenshot of the whole screen, a specific window, or a user-drawn region.
# Saves to ~/Pictures by default, but that can be changed via STATION_SCREENSHOT_DIR or XDG_PICTURES_DIR ENVs.
# Editor defaults to Satty but can be changed via --editor=<name> or STATION_SCREENSHOT_EDITOR env

[[ -f ~/.config/user-dirs.dirs ]] && source ~/.config/user-dirs.dirs
OUTPUT_DIR="${STATION_SCREENSHOT_DIR:-${XDG_PICTURES_DIR:-$HOME/Pictures}}"

if [[ ! -d $OUTPUT_DIR ]]; then
    notify-send "Screenshot directory does not exist: $OUTPUT_DIR" -u critical -t 3000
    exit 1
fi

pkill slurp && exit 0

SCREENSHOT_EDITOR="${STATION_SCREENSHOT_EDITOR:-satty}"

# Parse --editor flag from any position
ARGS=()
for arg in "$@"; do
    if [[ $arg == --editor=* ]]; then
        SCREENSHOT_EDITOR="${arg#--editor=}"
    else
        ARGS+=("$arg")
    fi
done
set -- "${ARGS[@]}"

open_editor() {
    local filepath="$1"
    if [[ $SCREENSHOT_EDITOR == "satty" ]]; then
        satty --filename "$filepath" \
            --output-filename "$filepath" \
            --actions-on-enter save-to-clipboard \
            --save-after-copy \
            --copy-command 'wl-copy'
    else
        $SCREENSHOT_EDITOR "$filepath"
    fi
}

MODE="${1:-smart}"
PROCESSING="${2:-slurp}"

# accounting for portrait/transformed displays
JQ_MONITOR_GEO='
def format_geo:
.x as $x | .y as $y |
    (.width / .scale | floor) as $w |
    (.height / .scale | floor) as $h |
    .transform as $t |
    if $t == 1 or $t == 3 then
        "\($x),\($y) \($h)x\($w)"
    else
        "\($x),\($y) \($w)x\($h)"
    end;
    '

    get_rectangles() {
        local active_workspace=$(hyprctl monitors -j | jq -r '.[] | select(.focused == true) | .activeWorkspace.id')
        hyprctl monitors -j | jq -r --arg ws "$active_workspace" "${JQ_MONITOR_GEO} .[] | select(.activeWorkspace.id == (\$ws | tonumber)) | format_geo"
        hyprctl clients -j | jq -r --arg ws "$active_workspace" '.[] | select(.workspace.id == ($ws | tonumber)) | "\(.at[0]),\(.at[1]) \(.size[0])x\(.size[1])"'
    }

# Select based on mode
case "$MODE" in
    region)
        hyprpicker -r -z >/dev/null 2>&1 & PID=$!
        sleep .1
        SELECTION=$(slurp 2>/dev/null)
        kill $PID 2>/dev/null
        ;;
    windows)
        hyprpicker -r -z >/dev/null 2>&1 & PID=$!
        sleep .1
        SELECTION=$(get_rectangles | slurp -r 2>/dev/null)
        kill $PID 2>/dev/null
        ;;
    fullscreen)
        SELECTION=$(hyprctl monitors -j | jq -r "${JQ_MONITOR_GEO} .[] | select(.focused == true) | format_geo")
        ;;
    smart|*)
        RECTS=$(get_rectangles)
        hyprpicker -r -z >/dev/null 2>&1 & PID=$!
        sleep .1
        SELECTION=$(echo "$RECTS" | slurp 2>/dev/null)
        kill $PID 2>/dev/null

  # If the selection area is L * W < 20, we'll assume you were trying to select whichever
  # window or output it was inside of to prevent accidental 2px snapshots
  if [[ $SELECTION =~ ^([0-9]+),([0-9]+)[[:space:]]([0-9]+)x([0-9]+)$ ]]; then
      if ((${BASH_REMATCH[3]} * ${BASH_REMATCH[4]} < 20)); then
          click_x="${BASH_REMATCH[1]}"
          click_y="${BASH_REMATCH[2]}"

          while IFS= read -r rect; do
              if [[ $rect =~ ^([0-9]+),([0-9]+)[[:space:]]([0-9]+)x([0-9]+) ]]; then
                  rect_x="${BASH_REMATCH[1]}"
                  rect_y="${BASH_REMATCH[2]}"
                  rect_width="${BASH_REMATCH[3]}"
                  rect_height="${BASH_REMATCH[4]}"

                  if ((click_x >= rect_x && click_x < rect_x + rect_width && click_y >= rect_y && click_y < rect_y + rect_height)); then
                      SELECTION="${rect_x},${rect_y} ${rect_width}x${rect_height}"
                      break
                  fi
              fi
          done <<<"$RECTS"
      fi
  fi
  ;;
esac

[[ -z $SELECTION ]] && exit 0

FILENAME="screenshot-$(date +'%Y-%m-%d_%H-%M-%S').png"
FILEPATH="$OUTPUT_DIR/$FILENAME"

if [[ $PROCESSING == "slurp" ]]; then
    grim -g "$SELECTION" "$FILEPATH" || exit 1
    wl-copy <"$FILEPATH"

    (
        ACTION=$(notify-send "Screenshot saved to clipboard and file" "Edit with Super + Alt + , (or click this)" -t 10000 -i "$FILEPATH" -A "default=edit")
        [[ $ACTION == "default" ]] && open_editor "$FILEPATH"
        ) &
    else
        grim -g "$SELECTION" - | wl-copy
fi