#!/bin/bash

###############################################################################
# Count Me In! v1.2 for Fatdog64 - A keystroke & mouse click counter		  #
# Copyright (C) JakeSFR 2018,2026											  #
# License: GNU GPL Version 2												  #
# Requirements: bash >= 4, libinput, yad, stdbuf, pr						  #
###############################################################################

export TEXTDOMAIN=countmein
export OUTPUT_CHARSET=UTF-8

APPNAME="Count Me In! v1.2"
DESCRIPTION="$(gettext "A keystroke & mouse click counter")"
COPYRIGHT="Copyright (C) JakeSFR 2018,2026"
LICENSE="License: GNU GPL Version 2"

# Globals
declare -A KEYS
declare -A CLICKS

TOTAL_KEYS=0
TOTAL_CLICKS=0
TOTAL_WHEEL=0

LAST_UPTIME=0
UPTIME=0
TIMESTAMP=$(date +%s)

KEYS_EXP=''
CLICKS_EXP=''

PERSISTENCE=false


# Constants
G_D="$(gettext "d")"	# days
G_H="$(gettext "h")"	# hours
G_M="$(gettext "m")"	# minutes
G_S="$(gettext "s")"	# seconds


TOOLTIP_HEAD="\
<span font-family='"'monospace'"'>\
<span font-size='"'x-large'"' background='"'gold'"'><b>$(gettext " *** ${APPNAME} *** ")</b></span>\n\
"

TOOLTIP_FOOT="\
</span>\
"


CONFDIR="${XDG_CONFIG_HOME}/countmein"
CONFFILE="${CONFDIR}/options"
CONFKEYS="${CONFDIR}/keys"
CONFCLICKS="${CONFDIR}/clicks"
CONFTOTALS="${CONFDIR}/totals"

WORKDIR="/tmp/countmein.${USER}.${$}"
ICON="/usr/share/pixmaps/themes/puppy16/mini-graph.xpm"
PIPE="${WORKDIR}/pipe"
RESET="${WORKDIR}/reset"

export APPNAME DESCRIPTION COPYRIGHT LICENSE ICON RESET PIPE CONFDIR CONFFILE PERSISTENCE

# -----------------------------------------------------------------------------

# Adapt dimensions to hi-dpi screens (Fatdog)
YAD_BIN="${YAD_BIN:-yad}"
YAD_VERSION="$($YAD_BIN --version 2>/dev/null)"

YAD_ABOUT_WIDTH=352
YAD_STATS_WIDTH=288
YAD_STATS_HEIGHT=512

if [ "${YAD_VERSION#*GTK+ 2}" != "$YAD_VERSION" -o "${GDK_SCALE:-1}" = 1 ] &&
	type fatdog-hidpi-lib >/dev/null 2>&1
then
	. fatdog-hidpi-lib
	set -- $(hidpi_apply_scale $YAD_ABOUT_WIDTH $YAD_STATS_WIDTH $YAD_STATS_HEIGHT)
	YAD_ABOUT_WIDTH=${1} YAD_STATS_WIDTH=${2} YAD_STATS_HEIGHT=${3}
fi

export YAD_BIN YAD_ABOUT_WIDTH YAD_STATS_WIDTH YAD_STATS_HEIGHT

# =============================================================================

run_as_root() {
	if [ $(id -u) -ne 0 ]; then
		exec gtksu "$APPNAME" "$0" "$@"
	fi
}

# -----------------------------------------------------------------------------

check_deps() {
	local dep allok=true

	if [ ! ${BASH_VERSINFO[0]} -ge 4 ]; then
		echo "$(gettext "Bash version is too old:") $BASH_VERSION"
		echo "$(gettext "Required version is >= 4")"
		allok=false
	fi

	for dep in libinput ${YAD_BIN} stdbuf pr; do
		if ! type ${dep} >/dev/null 2>&1; then
			echo "$(gettext "Missing dependency:") $dep"
			allok=false
		fi
	done

	if [ "${allok}" != 'true' ]; then
		echo "$(gettext "Aborting...")"
		exit 1
	fi
}

# -----------------------------------------------------------------------------

prepare() {
	mkdir -p "$WORKDIR"
	mkfifo "$PIPE"
	exec 3<>"$PIPE"

	trap "save_stats; pkill -P $$; exit" INT HUP TERM
	trap "rm -rf $WORKDIR" EXIT

	load_stats
}

# -----------------------------------------------------------------------------

save_stats() {
	mkdir -p "$CONFDIR"

	if [ -e "$CONFFILE" ]; then
		. "$CONFFILE"
	fi

	if [ "$PERSISTENCE" = "true" ]; then
		# Totals
		TOTAL_UPTIME=$(( ( $(date +%s) - ${TIMESTAMP} ) + ${LAST_UPTIME} ))
		echo -n "${TOTAL_KEYS}|${TOTAL_CLICKS}|${TOTAL_WHEEL}|${TOTAL_UPTIME}" > "$CONFTOTALS"

		# Keystrokes
		printf "%s\n" "${!KEYS[@]}" "${KEYS[@]}" | pr -l 999 -t2 -s'|' > "$CONFKEYS"

		# Mouse clicks
		printf "%s\n" "${!CLICKS[@]}" "${CLICKS[@]}" | pr -l 999 -t2 -s'|' > "$CONFCLICKS"
	fi
}

# -----------------------------------------------------------------------------

load_stats() {
	local key value

	if [ -e "$CONFFILE" ]; then
		. "$CONFFILE"
	fi

	if [ "$PERSISTENCE" = "true" ] && [ -e "$CONFTOTALS" ]; then
		IFS='|' read TOTAL_KEYS TOTAL_CLICKS TOTAL_WHEEL LAST_UPTIME < "$CONFTOTALS"
		
		if [ -e "$CONFKEYS" ]; then
			while IFS='|' read key value; do
				[ "$key" ] && [ "$value" ] && KEYS[$key]="${value}"
			done < "$CONFKEYS"
		fi

		if [ -e "$CONFCLICKS" ]; then
			while IFS='|' read key value; do
				[ "$key" ] && [ "$value" ] && CLICKS[$key]="${value}"
			done < "$CONFCLICKS"
		fi
	fi
}

# -----------------------------------------------------------------------------

about() {
	${YAD_BIN} --title="$(gettext "About")" \
		--window-icon="$ICON" \
		--button="gtk-ok:0" \
		--buttons-layout="center" \
		--width="${YAD_ABOUT_WIDTH}" \
		--margins="8" \
		--mouse \
		--text-align="center" \
		--text-justify="center" \
		--text="<b>${APPNAME}</b>" \
		--text-info <<< "$(echo -e "${DESCRIPTION}\n${COPYRIGHT}\n${LICENSE}")"
}

export -f about

# -----------------------------------------------------------------------------

statistics() {
	local wheel data_keys='' data_clicks='' data_wheel='' t_keys='' t_clicks='' t_wheel='' t_uptime=''

	declare -A s_keys
	declare -A s_clicks

	[ "${1#*KEYS}" ] && eval "s_keys="${1#*=}
	[ "${2#*CLICKS}" ] && eval "s_clicks="${2#*=}
	[ "$3" ] && wheel=${3} || wheel=0

	data_keys=$(printf "%s\n" "${!s_keys[@]}" "${s_keys[@]}" | pr -l 999 -t2 -s'|' | sort -rn -k2 -t '|')
	data_clicks=$(printf "%s\n" "${!s_clicks[@]}" "${s_clicks[@]}" | pr -l 999 -t2 -s'|' | sort -rn -k2 -t '|')
	[ ${wheel} -gt 0 ] && data_wheel="MOUSE_WHEEL|${wheel}" || data_wheel=''

	data_keys="${data_keys//KEY_/}"	# get rid of KEY_ prefix

	t_keys="$(printf " %-16s<b>%'-0d</b>" "$(gettext "Keystrokes:")" "$4")"
	t_clicks="$(printf " %-16s<b>%'-0d</b>" "$(gettext "Mouse Clicks:")" "$5")"
	t_wheel="$(printf " %-16s<b>%'-0d</b>" "$(gettext "Mouse Wheel:")" "$6")"
	t_uptime="$(printf " %-16s<b>%-0s</b>" "$(gettext "Uptime:")" "$7")"

	${YAD_BIN} --title="$(gettext "Statistics")" \
		--window-icon="$ICON" \
		--width="${YAD_STATS_WIDTH}" \
		--height="${YAD_STATS_HEIGHT}" \
		--mouse \
		--button="gtk-ok:0" \
		--buttons-layout="center" \
		--text "<span font-family='"'monospace'"'>${t_keys}\n${t_clicks}\n${t_wheel}\n${t_uptime}</span>" \
		--list \
		--no-selection \
		--dclick-action="" \
		--column="$(gettext "Key/Button"):TEXT" \
		--column="$(gettext "Strokes/Clicks"):NUM" \
		--sep-column="1" \
		--sep-value="-" \
		$(echo -e "${data_keys//|/\\n}\n-\n-\n${data_clicks//|/\\n}\n-\n-\n${data_wheel//|/\\n}")
}

export -f statistics

# -----------------------------------------------------------------------------

reset_stats() {
	${YAD_BIN} --title="$(gettext "Reset Statistics")" \
		--window-icon="$ICON" \
		--image="gtk-dialog-warning" \
		--mouse \
		--text-align="center" \
		--justify="center" \
		--text "<b>$(gettext "This action is irreversible!\nAre you really sure?")</b>"
	
	if [ $? -eq 0 ]; then
		touch "${RESET}"
	fi
}

export -f reset_stats

# -----------------------------------------------------------------------------

toggle_persistence() {
	local toggle

	mkdir -p "$CONFDIR"
	
	if [ -e "$CONFFILE" ]; then
		. "$CONFFILE"
	else
		echo "PERSISTENCE=$PERSISTENCE" > "$CONFFILE"
	fi
	
	if [ "$PERSISTENCE" = "false" ]; then
		toggle='true'
	else
		toggle='false'
	fi

	sed -i 's/PERSISTENCE=.*/PERSISTENCE='${toggle}'/' "$CONFFILE"
	
	update_menu "$toggle"
}

export -f toggle_persistence

# -----------------------------------------------------------------------------

update_menu() {
	local persistence_char
	
	case "$1" in
		'true')
			#PERSISTENCE_ICON='gtk-yes'
			persistence_char='✔'
		;;
		*)
			#PERSISTENCE_ICON='gtk-no'
			persistence_char='✘'
		;;
	esac

	cat > "$PIPE" << EOF
menu:\
$(gettext "Persistence")   ${persistence_char}!bash -c toggle_persistence!gtk-save|\
$(gettext "Reset stats")!bash -c reset_stats!gtk-refresh|\
|\
gtk-about!bash -c about|\
gtk-quit!quit
EOF
}

export -f update_menu

# -----------------------------------------------------------------------------

update_yad_tooltip() {
	local t_keys t_clicks t_wheel

	t_keys="$(printf "   %-20s%'-0d" "$(gettext "Keystrokes:")" $TOTAL_KEYS)"
	t_clicks="$(printf "   %-20s%'-0d" "$(gettext "Mouse Clicks:")" $TOTAL_CLICKS)"
	t_wheel="$(printf "   %-20s%'-0d" "$(gettext "Mouse Wheel:")" $TOTAL_WHEEL)"

	cat > "$PIPE" << EOF
tooltip: ${TOOLTIP_HEAD}\n\
<b>${t_keys}</b>\n\
<b>${t_clicks}</b>\n\
<b>${t_wheel}</b>\n\
\n\
   $(gettext "Uptime:") ${UPTIME}\n\
${TOOLTIP_FOOT}
EOF
}

# -----------------------------------------------------------------------------

update_yad_command() {
	cat > "$PIPE" << EOF
action: bash -c "statistics \"$KEYS_EXP\" \"$CLICKS_EXP\" \"$TOTAL_WHEEL\" \"$TOTAL_KEYS\" \"$TOTAL_CLICKS\" \"$TOTAL_WHEEL\" \"$UPTIME\""
EOF
}

# -----------------------------------------------------------------------------

create_tray() {
	${YAD_BIN} --kill-parent \
		--no-middle \
		--notification \
		--image="$ICON" \
		--listen <"$PIPE" &
	
	update_menu "$PERSISTENCE"
}

# -----------------------------------------------------------------------------

refresh_timeout() {
	while
		echo 'REFRESH'
	do
		sleep 1
	done
}

# -----------------------------------------------------------------------------

main_loop() {
	local null timeout type code changed='true'

	while
		case "$type" in

			'KEYBOARD_KEY')
				[ "$KEYS[$code]" ] && ((KEYS[$code]++)) || KEYS[$code]=1
				((TOTAL_KEYS++))
				changed='true'
			;;

			'POINTER_BUTTON')
				[ "$CLICKS[$code]" ] && ((CLICKS[$code]++)) || CLICKS[$code]=1
				((TOTAL_CLICKS++))
				changed='true'
			;;

			'POINTER_AXIS'|'POINTER_SCROLL_WHEEL')
				((TOTAL_WHEEL++))
				changed='true'
			;;

		esac

		# Refresh everything only on timeout
		if [ "$timeout" = 'REFRESH' ]; then
			UPTIME="$(eval echo $(date -ud @$(( $(date +%s) - $TIMESTAMP + $LAST_UPTIME )) +'$((%s/3600/24))${G_D} %H${G_H} %M${G_M} %S${G_S}'))"

			if [ -e "${RESET}" ]; then
				unset KEYS CLICKS
				declare -A KEYS CLICKS
				TOTAL_KEYS=0
				TOTAL_CLICKS=0
				TOTAL_WHEEL=0
				LAST_UPTIME=0
				UPTIME="0${G_D} 00${G_H} 00${G_M} 00${G_S}"
				TIMESTAMP=$(date +%s)
				changed='true'
				rm -f "$CONFKEYS" "$CONFCLICKS" "$CONFTOTALS"
				rm -f "$RESET"
			fi

			# Update tooltip
			update_yad_tooltip

			# Update command, but only if something's really changed
			if [ "$changed" = 'true' ]; then
				KEYS_EXP="$(declare -p KEYS)"
				CLICKS_EXP="$(declare -p CLICKS)"
				update_yad_command
				changed='false'
			fi
		fi
	do
		read timeout type null code null
	done < <( refresh_timeout & stdbuf -o0 libinput debug-events --show-keycodes | stdbuf -o0 grep -E 'released|(wheel)')
}

# -----------------------------------------------------------------------------

### START ###

run_as_root "$@"
check_deps
prepare
create_tray
main_loop

###############################################################################
