#!/usr/bin/env zsh

set -euo pipefail

readonly LOG_FILE="${HOME}/.screenshots.log"
readonly NOTIFY_TIME_LOW_MS=1500
readonly NOTIFY_TIME_NORMAL_MS=3000
readonly NOTIFY_TIME_CRITICAL_MS=5000
readonly CODE_CANCELED=22
readonly NOTIFY_LABEL="Screenshot"
readonly SCREENSHOTS_DIR="${SCREENSHOTS_DIR:-${HOME}/Pictures/screenshots}"
readonly TIMESTAMP_FORMAT="%H-%M-%S_%Y-%m-%d"
readonly NOOPENGL="${NOOPENGL:-}"
readonly NONOTIFY="${NONOTIFY:-}"


mkdir -p "${SCREENSHOTS_DIR}"


_notify() {
    local -r level="${1}"
    local -r notify_time_ms="${2}"
    local -r msg="${3}"

    if [ -n "${NONOTIFY}" ]; then
        return 0
    fi

    notify-send --urgency "${level}" \
                --expire-time "${notify_time_ms}" \
                "${NOTIFY_LABEL}" \
                "${msg}"
}

_notify_low() {
    _notify low "${NOTIFY_TIME_LOW_MS}" "${1}"
}

_notify_normal() {
    _notify normal "${NOTIFY_TIME_NORMAL_MS}" "${1}"
}


_notify_error() {
    _notify critical "${NOTIFY_TIME_CRITICAL_MS}" "Error: ${1}"
}

_capture_select() {
    local -r pipe="$(mktemp)"
    grim -g "$(slurp 2>"${pipe}")" -
    local -r code="$?"
    if grep -F "selection cancelled" <"${pipe}"; then
        rm "${pipe}"
        return "${CODE_CANCELED}"
    fi
    cat "${pipe}" 1>&2
    rm "${pipe}"
    return "${code}"
}

_capture_entire() {
    grim -
}

_capture_window() {
    grim -g "$(swaymsg -t get_tree | jq -r '.. | select(.type? == "con" and .focused == true) | .rect | "\(.x),\(.y) \(.width)x\(.height)"')" -
}

_copy_to_clipboard() {
    wl-copy -t image/png
}

_save_to_file() {
    cat > "${1}"
}

_process_screenshot() {
    local -r capture_func="${1}"
    local -r output_func="${2}"
    local -r success_msg="${3}"
    local -r error_msg="${4}"
    local -r output_arg="${5:-}"

    if ! ${capture_func} | ${output_func} "${output_arg}"; then
        if (( ${pipestatus[1]} == "${CODE_CANCELED}" )); then
            _notify_low "Selection was cancelled by keystroke."
            return 0
        fi
        _notify_error "${error_msg}"
        return 1
    fi

    _notify_normal "${success_msg}"
    return 0
}

_select_save() {
    local -r filename="${SCREENSHOTS_DIR}/selection_$(date +"${TIMESTAMP_FORMAT}").png"
    _process_screenshot _capture_select \
                        _save_to_file \
                        "Selection saved to ${filename}" \
                        "Failed to save selection screenshot" \
                        "${filename}"
}

_select_copy() {
    _process_screenshot _capture_select \
                       _copy_to_clipboard \
                       "Selection copied to clipboard" \
                       "Failed to copy selection to clipboard"
}

_window_save() {
    local -r filename="${SCREENSHOTS_DIR}/window_$(date +"${TIMESTAMP_FORMAT}").png"
    _process_screenshot _capture_window \
                       _save_to_file \
                       "Current window saved to ${filename}" \
                       "Failed to save Current window screenshot" \
                       "${filename}"
}

_window_copy() {
    _process_screenshot _capture_window \
                       _copy_to_clipboard \
                       "Current window copied to clipboard" \
                       "Failed to copy current window to clipboard"
}

_entire_save() {
    local -r filename="${SCREENSHOTS_DIR}/screen_$(date +"${TIMESTAMP_FORMAT}").png"
    _process_screenshot _capture_entire \
                       _save_to_file \
                       "Entire screen saved to ${filename}" \
                       "Failed to save entire screen" \
                       "${filename}"
}

_entire_copy() {
    _process_screenshot _capture_entire \
                       _copy_to_clipboard \
                       "Entire screen copied to clipboard" \
                       "Failed to copy entire screen to clipboard"
}

main() {
    if [ $# -ne 2 ]; then
        echo "Usage: ${0} {select|window|entire} {save|copy}"
        exit 1
    fi

    local -r target="${1}"
    local -r action="${2}"

    case "${target}-${action}" in
        select-save) _select_save &>> "${LOG_FILE}" ;;
        select-copy) _select_copy &>> "${LOG_FILE}" ;;
        window-save) _window_save &>> "${LOG_FILE}" ;;
        window-copy) _window_copy &>> "${LOG_FILE}" ;;
        entire-save) _entire_save &>> "${LOG_FILE}" ;;
        entire-copy) _entire_copy &>> "${LOG_FILE}" ;;
        *)
            echo "Invalid options: target=${target}, action=${action}"
            echo "Available targets: select, window, entire"
            echo "Available actions: save, copy"
            exit 1
            ;;
    esac
}

main "$@"