#!/usr/bin/env bash
#
# Copyright (C) 2024  Skylar Widulski
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>. 
#

VER='v1.5.5'

RUNTIME="${XDG_RUNTIME_DIR:-/run}/wf-shadow"
CONFIG="${WF_SHADOW_CONFIG:-${XDG_CONFIG_HOME:-$HOME/.config}/wf-shadow/config.sh}"

if [ ! -d "${CONFIG%/*}" ]; then
	mkdir -p "${CONFIG%/*}" 2>/dev/null
fi

if [ ! -f "$CONFIG" ]; then
	printf '' >> "$CONFIG"
fi

OUTPUT=""
AUDIO_DEV=""
AUDIO_CODEC=""
NO_AUDIO=""
VIDEO_DEV=""
VIDEO_CODEC=""
NO_VIDEO=""
VERBOSE=""
: "${WF_SHADOW_DIR:="$HOME"/Videos/wf-shadow}"
: "${WF_SHADOW_FMT:=%Y%m%d_%H:%M:%S}"
: "${WF_SHADOW_EXT:=mp4}"

. "$CONFIG"

if [ -n "$NO_AUDIO" ]; then
	NO_AUDIO=" "
fi

if [ -n "$NO_VIDEO" ]; then
	NO_VIDEO=" "
fi

if [ -n "$VERBOSE" ]; then
	VERBOSE=" "
fi

trap 'kill $(jobs -p) 2>/dev/null' EXIT
trap 'kill "$(<"$RUNTIME"/"$OUTPUT"_host.pid)"; rm -f "$RUNTIME"/"$OUTPUT"_host.pid' INT
set -m
shopt -s extglob

if [ ! -d "$RUNTIME" ]; then
	mkdir -p "$RUNTIME" 2>/dev/null
fi

if [ ! -d /tmp/wf-shadow ]; then
	mkdir -p /tmp/wf-shadow 2>/dev/null
fi

if [ ! -d "$WF_SHADOW_DIR" ]; then
	mkdir -p "$WF_SHADOW_DIR" 2>/dev/null
fi

if ! command -v wf-recorder &>/dev/null; then
	printf "Need wf-recorder\n"
	exit 1
fi

if ! command -v wlr-randr &>/dev/null; then
	printf "Need wlr-randr\n"
	exit 1
fi

print_help() {
	cat << EOF
wf-shadow: A Wayland shadow recorder using wf-recorder and wofi

USAGE: wf-shadow OPTION [OPTIONS...]

OPTIONS:
	-s, --start [OUTPUT]	Start a wf-shadow instance with output number OUTPUT.
	-r, --record [OUTPUT]	Record the last few moments of video using a wofi menu.
	-e, --end [OUTPUT]	End wf-shadow instance with output number OUTPUT.
	-a, --audio DEVICE	Use DEVICE for audio recording. Passed to wf-recorder -a.
	-V, --video DEVICE	Use DEVICE for video encoding. Passed to wf-recorder -d.
	-A, --acodec CODEC	Use CODEC as the audio codec in ffmpeg encoding.
	-c, --vcodec CODEC	Use CODEC as the video codec in ffmpeg encoding.
	-n, --no-audio	Disable audio recording. Overrides audio options.
	-N, --no-video	Disable video recording. Overrides video options.
	-y, --yes-audio	Enable audio recording. Overrides -n.
	-Y, --yes-video	Enable video recording. Overrides -N.
	-R, --raw-record OUTPUT FILE SECONDS	Low level record. Passed directly to the \`record\` function.
	-v, --verbose	Verbose output.
	-h, --help	Show this menu.
	--version	Display version information.

For a more in-depth guide, see wf-shadow(1)
EOF
	exit $1
}

record() {
	local SAVE_FILE
	if [ ! $WF_SHADOW_SAVE ]; then
		SAVE_FILE="$WF_SHADOW_DIR"/"$(date +"$WF_SHADOW_FMT")__$2"."$WF_SHADOW_EXT"
	else
		SAVE_FILE="$(eval "$WF_SHADOW_SAVE")"
	fi
	if eval "ffmpeg -nostdin -y -sseof -'$2' -i '$1' ${NO_VIDEO:-${VIDEO_CODEC:+-vcodec "$VIDEO_CODEC"}} ${NO_AUDIO:-${AUDIO_CODEC:+-acodec "$AUDIO_CODEC"}} ${NO_VIDEO:+-vn} ${NO_AUDIO:+-an} '$SAVE_FILE' ${VERBOSE:-&>/dev/null}"; then
		notify-send "Clip saved" "$2 second clip saved to $SAVE_FILE"
	else
		notify-send "Error saving clip" "Start a new session with --verbose to see ffmpeg output"
	fi
}

start() {
	while true; do
		for i in /tmp/wf-shadow/"$OUTPUT"-{00..10}.mp4; do
			printf '' >> "$i"
			printf '%s' "$i" > "$RUNTIME"/"$OUTPUT"_current
			if [ -n "$VERBOSE" ]; then
				wf-recorder ${OUTPUT:+-o "$OUTPUT"} ${VIDEO_DEV:+-d "$VIDEO_DEV"} -a $AUDIO_DEV -f "$i" <<< y$'\n'1 &
			else
				wf-recorder ${OUTPUT:+-o "$OUTPUT"} ${VIDEO_DEV:+-d "$VIDEO_DEV"} -a $AUDIO_DEV -f "$i" <<< y$'\n'1 &>/dev/null &
			fi
			printf "%s" "$!" > "$RUNTIME"/"$OUTPUT".pid
			fg || exit
		done
	done
}

check_output() {
	if ! wlr-randr | grep "$1" &>/dev/null; then
		printf "Output %s does not exist\n" "$1"
		exit 1
	fi	
}

until [ -z "$1" ]; do
	re="^[A-Z]+-[1-9][0-9]*$"
	case "$1" in
		'-h' | '--help') print_help ;;
		'-e' | '--end')
			END=1
			if [[ "$2" =~ $re ]]; then
				check_output "$2"
				OUTPUT="$2"
				shift 2
			else
				shift 1
			fi ;;
		'-r' | '--record')
			RECORD=1
			if [[ "$2" =~ $re ]]; then
				check_output "$2"
				OUTPUT="$2"
				shift 2
			else
				shift 1
			fi ;;
		'-s' | '--start')
			START=1
			if [[ "$2" =~ $re ]]; then
				check_output "$2"
				OUTPUT="$2"
				shift 2
			else
				shift 1
			fi ;;
		'-a' | '--audio')
			AUDIO_DEV="$2"
			shift 2 ;;
		'-A' | '--acodec')
			if [ -n "$2" ]; then
				AUDIO_CODEC="$2"
				shift 2
			else
				print_help
			fi ;;
		'-n' | '--no-audio')
			NO_AUDIO=" "
			shift 1 ;;
		'-y' | '--yes-audio')
			NO_AUDIO=""
			shift 1 ;;
		'-V' | '--video')
			VIDEO_DEV="$2"
			shift 2 ;;
		'-c' | '--vcodec')
			if [ -n "$2" ]; then
				VIDEO_CODEC="$2"
				shift 2
			else
				print_help
			fi ;;
		'-N' | '--no-video')
			NO_VIDEO=" "
			shift 1 ;;
		'-Y' | '--yes-video')
			NO_VIDEO=""
			shift 1 ;;
		'-R' | '--raw-record')
			if [ -n "$4" ]; then
				OUTPUT="$2"
				FILE="$3"
				TIME="$4"
				RAW_RECORD=1
				shift 4
			else
				print_help
			fi ;;
		'-d' | '--debug')
			set -x
			shift 1 ;;
		'-v' | '--verbose')
			VERBOSE=" "
			shift 1 ;;
		'--version')
			printf "wf-shadow %s\n" "$VER"
			wf-recorder -v
			if command -v wofi &>/dev/null; then
				printf "wofi %s\n" "$(wofi -v)"
			fi
			exit ;;
		*) print_help ;;
	esac
done


if [ -n "$NO_VIDEO" ] && [ -n "$NO_AUDIO" ]; then
	printf "Must have either audio or video recording enabled.\n"
	exit 3
fi

if [ "$START" == 1 ]; then
	if [ ! -e "$RUNTIME"/"$OUTPUT"_host.pid ]; then
		printf "%s" "$$" > "$RUNTIME"/"$OUTPUT"_host.pid
		start
	else
		printf "A wf-shadow instance is already running on output %s (%s/%s_host.pid)\n" "$OUTPUT" "$RUNTIME" "$OUTPUT"
	fi
fi

if [ "$RAW_RECORD" == 1 ]; then
	kill -2 "$(<"$RUNTIME"/"$OUTPUT".pid)"
	record "$FILE" "$TIME"

fi

if [ "$RECORD" == 1 ]; then
	if ! command -v wofi &>/dev/null; then
		printf "Need wofi\n"
		exit 1
	fi
	if [ -e "$RUNTIME"/"$OUTPUT"_host.pid ]; then
		WOFI_STRING="5 seconds\n10 seconds\n15 seconds\n30 seconds\n1 minute\n5 minutes\n10 minutes\nCancel\n"
		TIME="$(eval "printf '$WOFI_STRING' | wofi --dmenu -i -H 500 -W 250 -x 0 -y 0 ${VERBOSE:-2>/dev/null}")"
		FILE="$(<"$RUNTIME"/"$OUTPUT"_current)"
		kill -2 "$(<"$RUNTIME"/"$OUTPUT".pid)"
		case "$TIME" in
			'5 seconds')  record "$FILE" 5   ;;
			'10 seconds') record "$FILE" 10  ;;
			'15 seconds') record "$FILE" 15  ;;
			'30 seconds') record "$FILE" 30  ;;
			'1 minute')   record "$FILE" 60  ;;
			'5 minutes')  record "$FILE" 300 ;;
			'10 minutes') record "$FILE" 600 ;;
			+([0-9]))     record "$FILE" "$TIME" ;;
		esac
	else
		printf "No such session on output %s\n" "$OUTPUT"
	fi
fi

if [ "$END" == 1 ]; then
	if [ -e "$RUNTIME"/"$OUTPUT"_host.pid ]; then
		kill "$(<"$RUNTIME"/"$OUTPUT"_host.pid)"
		rm -f "$RUNTIME"/"$OUTPUT"_host.pid
	else
		printf "No such session on output %s\n" "$OUTPUT"
	fi
fi