#!/bin/env bash

# CLI options
DMENU_CMD=""
while getopts ":d:" opt; do
    case "$opt" in
        d) DMENU_CMD="$OPTARG" ;;
        *) ;;
    esac
done
shift $((OPTIND - 1))

if [ -z "$DMENU_CMD" ]; then
    notify-send "Missing launcher" "Use -d to provide a dmenu command"
    exit 1
fi
read -r -a DMENU_ARR <<< "$DMENU_CMD"
DMENU_BIN="${DMENU_ARR[0]}"
command -v "$DMENU_BIN" &>/dev/null || { notify-send "Launcher not found" "$DMENU_BIN is not installed"; exit 1; }

# Directories
SCREENSHOT_DIR="${HOME}/Pictures"
RECORDING_DIR="${HOME}/Videos"

# Ensure directories exist
mkdir -p "$SCREENSHOT_DIR" "$RECORDING_DIR"

# Stop active wl-screenrec session
if pgrep -u "$USER" wl-screenrec > /dev/null; then
	pkill -INT -u "$USER" wl-screenrec
	notify-send "Recording stopped"
	
	# Ask if user wants to compress the video
	COMPRESS=$(printf "%b" "✅ compress\n❌ keep original" | "${DMENU_ARR[@]}" "🗜️ Compress video?")
	if [[ "$COMPRESS" == "✅ compress" ]]; then
		# Find the most recent video file
		LATEST_VIDEO=$(find "$RECORDING_DIR" -name "*.mp4" -type f -printf '%T@ %p\n' | sort -n | tail -1 | cut -d' ' -f2-)
		if [[ -n "$LATEST_VIDEO" ]] && [[ -f "$LATEST_VIDEO" ]]; then
			COMPRESSED_VIDEO="${LATEST_VIDEO%.mp4}_compressed.mp4"
			notify-send "Compressing video..." "This may take a moment"
			if ffmpeg -i "$LATEST_VIDEO" -c:v libx264 -crf 23 -preset medium -c:a aac -b:a 128k "$COMPRESSED_VIDEO" -y &>/dev/null; then
				notify-send "Compression complete" "$COMPRESSED_VIDEO"
				# Ask if user wants to delete original
				DELETE_ORIGINAL=$(printf "%b" "🗑️ delete original\n📁 keep both" | "${DMENU_ARR[@]}" "Delete original?")
				if [[ "$DELETE_ORIGINAL" == "🗑️ delete original" ]]; then
					rm "$LATEST_VIDEO"
					notify-send "Original deleted" "Keeping compressed version only"
				fi
			else
				notify-send "Compression failed" "Keeping original file"
			fi
		fi
	fi
	exit 0
fi

# Get current outputs dynamically
OUTPUTS=$(hyprctl monitors | grep "Monitor" | awk '{print $2}' | paste -sd " " -)

TOP_LEVEL=$(printf "%b" "screenshot\nrecord" | "${DMENU_ARR[@]}" "󰄀 ")
[ -z "$TOP_LEVEL" ] && exit 0

if [ "$TOP_LEVEL" = "screenshot" ]; then
	OPTIONS=$(
		cat <<EOF
screenshot selection
$(for OUTPUT in $OUTPUTS; do echo "screenshot $OUTPUT"; done)
screenshot all
EOF
	)
else
	OPTIONS=$(
		cat <<EOF
record selection
$(for OUTPUT in $OUTPUTS; do echo "record $OUTPUT"; done)
record all
EOF
	)
fi

SELECTION=$(printf "%b" "$OPTIONS" | "${DMENU_ARR[@]}" "󰄀 ")
[ -z "$SELECTION" ] && exit 0

# Ask for audio mode if it's a recording
select_audio_mode() {
    AUDIO_MODE=$(printf "%b" "🎤 mic\n💻 internal\n🚫 none" | "${DMENU_ARR[@]}" "🎙️ Audio?")
    [ -z "$AUDIO_MODE" ] && echo "__cancel__" && return

    case "$AUDIO_MODE" in
        "🚫 none")
            echo ""
            ;;
        "🎤 mic")
            echo "--audio"
            ;;
        "💻 internal")
            echo "--audio --audio-device default.monitor"
            ;;
        *)
            echo ""
            ;;
    esac
}

open_in_satty() {
	[ -x "$(command -v satty)" ] || return 0
	satty -f "$1" >/dev/null 2>&1 &
}

# Screenshot or record depending on selection
case "$SELECTION" in
	"screenshot selection")
		TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)"
		IMG="${SCREENSHOT_DIR}/${TIMESTAMP}.png"
		grim -g "$(slurp)" "$IMG" || { notify-send "Error" "Failed to take screenshot"; exit 1; }
		[ -x "$(command -v wl-copy)" ] && wl-copy < "$IMG"
		open_in_satty "$IMG"
		notify-send "Screenshot Taken" "$IMG"
		;;

	"screenshot all")
		TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)"
		IMG="${SCREENSHOT_DIR}/${TIMESTAMP}.png"
		TEMP_FILES=()
		for OUTPUT in $OUTPUTS; do
			OUT_IMG="${IMG//.png/-$OUTPUT.png}"
			grim -c -o "$OUTPUT" "$OUT_IMG" && TEMP_FILES+=("$OUT_IMG")
		done
		if [ "${#TEMP_FILES[@]}" -gt 1 ]; then
			montage "${TEMP_FILES[@]}" -tile x1 -geometry +0+0 "$IMG"
			rm "${TEMP_FILES[@]}"
			[ -x "$(command -v wl-copy)" ] && wl-copy < "$IMG"
			open_in_satty "$IMG"
			notify-send "Screenshot Taken" "$IMG"
		else
			notify-send "Error" "Not enough outputs to montage"
		fi
		;;

	"record selection")
		TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)"
		VID="${RECORDING_DIR}/${TIMESTAMP}.mp4"
		AUDIO_ARGS=$(select_audio_mode)
		[ "$AUDIO_ARGS" == "__cancel__" ] && notify-send "Recording cancelled" && exit 0
		
		# Get selection geometry
		notify-send "Select area" "Click and drag to select recording area"
		GEOMETRY="$(slurp)"
		if [ $? -ne 0 ] || [ -z "$GEOMETRY" ]; then
			notify-send "Recording cancelled" "Area selection was cancelled"
			exit 0
		fi
		
		# Start recording
		if [ -n "$AUDIO_ARGS" ]; then
			eval "wl-screenrec --geometry '$GEOMETRY' --filename '$VID' $AUDIO_ARGS" &
		else
			wl-screenrec --geometry "$GEOMETRY" --filename "$VID" &
		fi
		notify-send "Recording started" "Saving to $VID"
		;;

	"record all")
		notify-send "Not implemented" "Recording all screens is not supported yet."
		;;

	"screenshot "*)
		TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)"
		IMG="${SCREENSHOT_DIR}/${TIMESTAMP}.png"
		OUT=$(echo "$SELECTION" | awk '{print $2}')
		grim -c -o "$OUT" "$IMG" || { notify-send "Error" "Failed to screenshot $OUT"; exit 1; }
		[ -x "$(command -v wl-copy)" ] && wl-copy < "$IMG"
		open_in_satty "$IMG"
		notify-send "Screenshot Taken" "$IMG"
		;;

	"record "*)
		TIMESTAMP="$(date +%Y-%m-%d_%H-%M-%S)"
		VID="${RECORDING_DIR}/${TIMESTAMP}.mp4"
		OUT=$(echo "$SELECTION" | awk '{print $2}')
		AUDIO_ARGS=$(select_audio_mode)
		[ "$AUDIO_ARGS" == "__cancel__" ] && notify-send "Recording cancelled" && exit 0
		
		# Start recording
		if [ -n "$AUDIO_ARGS" ]; then
			eval "wl-screenrec --output '$OUT' --filename '$VID' $AUDIO_ARGS" &
		else
			wl-screenrec --output "$OUT" --filename "$VID" &
		fi
		notify-send "Recording started" "Saving to $VID"
		;;

	*)
		notify-send "Invalid selection" "$SELECTION"
		;;
esac