wmpersist

#!/bin/bash
#
# wmpersist
#   - save/load browser window titles to winname mappings
#   - works with ratpoison window manager
#
# usage:
#   "wmpersist save X": write out chromium window titles to
#                     ratpoison winname mappings
#
#   "wmpersist load X": take existing windows, search for
#                     titles in the save file, and
#                     rename/renumber within ratpoison to
#                     match them, once chromium has been
#                     restarted (typically following a crash,
#                     which happens frequently)
#
#   second arg (X): chrome, brave or nightly (browser to save/load for)
#
# scott@smemsh.net
# https://github.com/smemsh/utilx/
# https://spdx.org/licenses/GPL-2.0
#
##############################################################################

invname=${0##*/}
setenv () { local v=$1; shift; IFS= eval $v="\$*"; }
warn   () { echo "$invname: ${FUNCNAME[1]}: $*" >&2; }
bomb   () { warn "$@"; false; exit; }

# see get_window_titles()
declare -A classname1 classname2
classname1[chrome]='chromium'
classname2[chrome]='Chromium'
classname1[brave]='brave-browser'
classname2[brave]='Brave-browser'
classname1[nightly]='brave-browser-nightly'
classname2[nightly]='Brave-browser-nightly'

# enable for additional prints, etc
debug=${DEBUG:-0}

# while transitioning ratpoison -> sdorfehs, we should work under either
[[ $SDORFEHS_PID || $(wmname) == sdorfehs ]] && wm=sdorfehs || wm=ratpoison
declare -A cmd_group_list cmd_group_move cmd_group_select
cmd_group_list[sdorfehs]=vscreens
cmd_group_list[ratpoison]=groups
cmd_group_move[sdorfehs]=vmove
cmd_group_move[ratpoison]=gmove
cmd_group_select[sdorfehs]=vselect
cmd_group_select[ratpoison]=gselect

if ((debug))
then
	debug_file=~/var/$wm/wmpersist.debug
	trunc $debug_file
fi

rpcmd="$wm -c"

# total number of chromium windows we loop about
#
declare -g nwins

# array for storing the current names, numbers, titles
#
declare -a windata_names
declare -a windata_groups
declare -a windata_numbers
declare -a windata_titles

# dictionary to save win nums, check if exists upon restore
#
declare -A windata_number_to_title

# dictionaries for mappings against title, so we can take
# the current one and make it the desired (saved) one
#
declare -A winmap_title_to_name
declare -A winmap_title_to_number

###

# for switching on invocation name based on whether a
# correspondingly named function exists
#
function_exists ()
{
	declare -f $1 &>/dev/null
}

###

# use ratpoison 'windows' command to construct lines with format:
# "xid winname number" where: xid is the xwindows id, winname
# is the res_name class hint, and number is the ratpoison
# window number (so we can order them on reconstruct), and
# write to stdout
#
get_window_rpdata ()
{
	local i curgroup group groupname
	local -a groups groupnames

	# TODO share with get_window_titles(), needs change there too
	local sep="^\\w+\\s+\\d+\\s${classname1}\\s${classname2}\\s"

	groups=($($rpcmd ${cmd_group_list[$wm]}))
	for ((i = 0; i < ${#groups[@]}; i++)); do
		[[ ${groups[i]} =~ [[:digit:]]+(.)(.*) ]]
		groupname=${BASH_REMATCH[2]}
		groupnames+=($groupname)
		[[ ${BASH_REMATCH[1]} == '*' ]] &&
			curgroup=$groupname
	done

	(for group in ${groupnames[@]}; do
		$rpcmd "${cmd_group_select[$wm]} $group"
		$rpcmd "windows $group %i %a %c %t %n" |
		grep -P "$sep" |
		awk '{
			printf("%s %u", $1, $2)
			for (i = 5; i <= NF; i++)
				printf(" %s", $i)
			printf("\n")
		}'
	done) | sort -snrk 2,2

	$rpcmd "${cmd_group_select[$wm]} $curgroup"
}

# use 'xwininfo' to print a tree of windows and fiter out
# fields with just xid and title separated by a single space
# and written to lines on stdout
#
get_window_titles ()
{
	# TODO share with get_window_rpdata(), needs change there too
	local sep="(\"$classname1\" \"$classname2\")"

	xwininfo -int -root -children |
	grep -P '^\s+\d+\s' |
	grep ":.*$sep" |
	awk -F "$sep" '{print $1}' |
	sed s,^'[[:space:]]*',, |
	sed -r 's,^([[:digit:]]+)[[:blank:]]+"(.*)":.\(,\1 \2,' |
	sort -nrk 1,2
}

# see https://tronche.com/gui/x/icccm/sec-4.html for
# description of withdrawn window; it may be that ratpoison
# should do something with these windows and doesn't; unknown,
# but chrome leaves these around sometimes, seem to be eg
# contacts.google.com, possibly bookmarks window (latter
# unverified, but that would explain an earlier problem)
#
# TODO: more study needed on what these withdrawn windows are
# and in what circumstances they arise and why, also is
# ratpoison supposed to do something to destroy them?
#
window_is_withdrawn ()
{
	# this magic incantation gets the WM_STATE.state
	# field, which is zeroed if the window is withdrawn
	#
	local wmstate=$(
		xprop \
			-notype \
			-id $1 \
			-f WM_STATE 32c ' $0' \
			WM_STATE \
		| awk '{print $NF}'
	)

	# ICCCM 4.1.3.1: "Top-level windows in the Withdrawn
	# state may or may not have the WM_STATE property ...
	# When the window is withdrawn, the window manager will
	# either change the state field's value to
	# WithdrawnState [ie 0] or it will remove the WM_STATE
	# property entirely." (note: only saw 0 values for
	# withdrawn windows prior to 20170202 -- chromium
	# windows always had a WM_STATE -- but now we're
	# starting to see some without it, hence the new
	# condition)
	#
	if [[ ! $wmstate =~ ^[[:digit:]]$ ]] || ((wmstate == 0))
	then true
	else false
	fi
}

# create variables used by save and restore from ratpoison
# (xid, name, number) and xwininfo (xid, title) output
#
gather_data ()
{
	local -a rpdata titles
	local i

	# whether to check that the rp and xwininfo lists
	# match (they should on a save, they won't on a restore)
	#
	local should_match=$1

	# for storing the interesting variables from each
	# "xwininfo" and "ratpoison -c windows" lines, which
	# we will process later
	#
	local rpxid rpnam rpgrp rpnum
	local ttxid title
	
	# pull out the ratpoison data fields we want
	# from each output line stored in the array
	#
	gather_rpdata_fields ()
	{
		setenv pattern \
			'^([^[:space:]]+)[[:space:]]+' \
			'([[:digit:]]+)[[:space:]]' \
			'(.*)[[:space:]]' \
			'([[:digit:]]+)' \
		;
		[[ "${rpdata[i]}" =~ $pattern ]]
		rpgrp=${BASH_REMATCH[1]}
		rpxid=${BASH_REMATCH[2]}
		rpnam=${BASH_REMATCH[3]}
		rpnum=${BASH_REMATCH[4]}
	}

	# and then also save fields we want from the xwininfo
	# output, taking care to ignore withdrawn windows
	#
	gather_xwininfo_fields ()
	{
		# keep iterating until we find one not withdrawn, in
		# usual case this loop runs only once
		#
		while [[ ${titles[i]} ]]
		do
			setenv pattern \
				'([[:digit:]]+)[[:space:]]' \
				'(.*)' \
			;
			[[ "${titles[i]}" =~ $pattern ]]
			ttxid=${BASH_REMATCH[1]}
			title=${BASH_REMATCH[2]}

			# ratpoison will not list withdrawn windows,
			# which chrome seems to leave around sometimes
			# (causing a mismatch or one list longer than
			# the other if last element) so check if this is
			# the case
			#
			if window_is_withdrawn $ttxid
			then
				# withdrawn windows can just be ignored,
				# since they are not visible to the user
				# anyways
				#
				((debug)) && echo "ignoring $ttxid"

				# however, now we need to rebuild the
				# $titles[] array without this bogus
				# entry
				#
				unset "titles[i]"
				titles=("${titles[@]}")

				# restart at top to try again for
				# a match that isn't withdrawn; again,
				# this "loop" usually runs only once
				# (this is the only looping point)
				#
				continue
			fi

			# on save, the two arrays $rpdata[] and
			# $titles[] should have the same windows in the
			# same order with different data but the first
			# datum identical (the xid), so we can now merge
			# the lines, but compare their xids (the join
			# field) to be sure they are the same: a small
			# race window exists between collecting the
			# window data from ratpoison, and from xwininfo,
			# when a new window could have been created
			#
			if ((should_match && rpxid != ttxid))
			then
				# we cannot explain the mismatch as a
				# withdrawn window (already tested for)
				# so probably a new window was created
				# during race window (unhandled; abort)
				#
				bomb \
					"winid mismatch" \
					"rpdata:$rpxid" \
					"titles:$ttxid" \
				;
			else
				# either this is a restore and
				# $should_match is false, or we have a
				# good pattern match; either way, done
				#
				return
			fi
		done

		# we ran out of withdrawn windows to shift without
		# finding the one corresponding to $rpdata[i]
		#
		bomb "shifted $titles[] until no more, without finding match"
	}

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

	# ratpoison does not store the window title accessible
	# from 'windows' command so we have to do this twice:
	# once for ratpoison, and once to get the titles,
	# making the $rpdata[] and $titles[] arrays which have
	# one element per line (we split up later into the
	# globals declared above for each iteration, and run
	# the callback)
	#
	readarray -t rpdata <<< "$(get_window_rpdata)"
	readarray -t titles <<< "$(get_window_titles)"

	# ratpoison will govern how many "real" windows we
	# have; the xwininfo data has some extraneous ones
	# which we will need to filter out later
	#
	nwins=${#rpdata[@]}

	if ((debug))
	then
		declare -p rpdata >> $debug_file
		declare -p titles >> $debug_file
		declare -p nwins >> $debug_file
	fi

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

	for ((i = 0; i < nwins; i++))
	do
		gather_rpdata_fields
		gather_xwininfo_fields

		# stash the gathered variables in indexed (by
		# line number from $rpdata[] and $titles[],
		# which should be same) and associative arrays
		# (by title) to use later (the save/restore
		# routines both use these data)
		#
		windata_names+=("$rpnam")
		windata_groups+=("$rpgrp")
		windata_titles+=("$title")
		windata_numbers+=($rpnum)

		# also maintain a dictionary of rpnums: if
		# restore routine has called us, it will want
		# to look up whether the saved number it's
		# trying to restore already exists, and also
		# will want to swap mappings, so we use
		# this to mark that the number has an
		# associated title already, and what the title
		# is (so we can swap our own list in-place)
		#
		windata_number_to_title[$rpgrp,$rpnum]="$title"

		# however, we *do* want to use a title:number
		# dictionary because, on restore, we have to
		# know existing mapping of title to ratpoison
		# window number (we could just search the
		# array and look up name at the corresponding
		# offset of title, but it would happen for
		# every window we wanted to change the name
		# of, and that's too slow); this way we can
		# search the array in constant time for each
		# one.  when found we get the number, switch
		# to it, and then write the name
		#
		winmap_title_to_number["$title"]="$rpgrp,$rpnum"

		# we don't actually use this yet, because we
		# already have the existing mappings, gathered
		# in $windata_{names,titles}[] and only need
		# it for creating the save-state file, which
		# means no searching is required (we just
		# iterate over the elements and write each
		# one to the save file)
		#
		winmap_title_to_name["$title"]="$rpnam"
	done

	if ((debug))
	then
		declare -p windata_names >> $debug_file
		declare -p windata_groups >> $debug_file
		declare -p windata_numbers >> $debug_file
		declare -p windata_titles >> $debug_file
		declare -p winmap_title_to_name >> $debug_file
		declare -p winmap_title_to_number >> $debug_file
	fi

	# any remaining windows in $titles[] are not known to
	# ratpoison: either withdrawn -- and therefore can be
	# ignored -- or new, in which case we also just ignore
	# them (wouldn't be windows for us anyways since we the
	# operator aren't trying to race our script against new
	# browser windows), so we can safely exit now
}

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

# paste together the stashed fields from $rpdata[] and
# $titles[] into now-merged lines with the info we need to
# reconstruct, and write them to the save file
#
# TODO needs some kind of locking, especially if we will run
# this on all window title change events as we plan to via
# hook NOTE do not see title-change hook in manual, but is in
# source code?!?!
#
save    () { store "$@"; }
dump    () { store "$@"; }
store   ()
{
	local i

	mv -i $save_file $save_file.`now`
	gather_data 1

	# write the relevant stashed window data (as gathered
	# by gather_data()) into the designated save file
	#
	for ((i = 0; i < nwins; i++)); do
		printf "%s %u %s %s\n" \
			"${windata_groups[i]}" \
			"${windata_numbers[i]}" \
			"${windata_names[i]}" \
			"${windata_titles[i]}" \
		>> $save_file
	done

	echo "saved" # TODO ratpoison verbexec does not display this?
}

restore () { load "$@"; }
reload  () { load "$@"; }
load    ()
{
	local i
	local pattern line curgroup
	local saved_grp saved_num saved_name saved_title
	local target target_num target_grp target_cmp saved_cmp

	local -a lines groups

	test -e $save_file ||
		bomb "persist data file does not exist"

	readarray -t lines < $save_file

	if ((debug)); then
		declare -p lines > $debug_file; fi

	gather_data 0

	print_rpinfo () { $rpcmd "info %i %t %n"; }

	# for each saved title-name mapping, find the target
	# window, switch to it, rename it (in ratpoison, ie
	# 'winname' set by 'title'), and switch back
	#
	for ((i = 0; i < ${#lines[@]}; i++))
	do
		line="${lines[i]}"
		setenv pattern '^' \
			'^([^[:space:]]+)[[:space:]]+' \
			'([[:digit:]]+)[[:space:]]' \
			'([^[:space:]]+)[[:space:]]' \
			'(.*)' \
		;
		[[ $line =~ $pattern ]]
		saved_grp=${BASH_REMATCH[1]}
		saved_num=${BASH_REMATCH[2]}
		saved_name=${BASH_REMATCH[3]}
		saved_title=${BASH_REMATCH[4]}

		target=${winmap_title_to_number["$saved_title"]}
		target_num=${target#*,}
		target_grp=${target%,*}
		target_cmp=$target_grp,$target_num
		saved_cmp=$saved_grp,$saved_num
		if ! [[ $target_num && $target_grp ]]; then
			bomb "failed map title -> target: \"$saved_title\""; fi

		if ((debug))
		then
			printf "%s: %s\n" \
		             loop "$i" \
		        saved_grp "$saved_grp" \
		        saved_cmp "$saved_cmp" \
		       saved_name "$saved_name" \
		      saved_title "$saved_title" \
		       target_cmp "$target_cmp" \
		; fi

		if ((debug)); then
			printf "before: %s\n" `print_rpinfo`; fi

		# TODO: instead of having to select the
		# window, there should be a way to give a
		# second arg for target window, as with the
		# 'number' ratpoison command
		#
		$rpcmd "${cmd_group_select[$wm]} $target_grp" ||
			bomb "failed switch to target group $target_grp"
		$rpcmd "select $target_num" ||
			bomb "failed to select target window $target_num"

		if ((debug)); then
			printf "after-select: %s\n" `print_rpinfo`; fi

		# sets rp $winname (res_name class hint for us)
		#
		# TODO this doesn't seem to actually set
		# res_name but some internal ratpoison
		# variable, which is the reason for a whole
		# lot of trickery in this script, otherwise we
		# could just set the xprop directly, out of
		# band, without using the window manager at
		# all
		#
		$rpcmd "title $saved_name" ||
			bomb "failed to set title $saved_name"

		if ((debug)); then
			printf "after-title: %s\n" `print_rpinfo`; fi

		# if target window starts out in different group than
		# the destination (ie saved group), we have to move it
		#
		if [[ $target_grp != $saved_grp ]]
		then
			$rpcmd "${cmd_group_move[$wm]} $saved_grp" ||
				bomb "could not move window to group $saved_grp"

			# window move does not change current group in
			# ratpoison, but does in sdorfehs.  this should be
			# idempotent anyways
			#
			$rpcmd "${cmd_group_select[$wm]} $saved_grp"

			# we have a switchgroup hook that always runs
			# "other" so we have to switch back to the target
			# window.  if the group had no windows, the "other"
			# will fail because there's "no other window" but
			# also will this second "other" fail and so we'll
			# still end up on the right window (the only one in
			# the group, if the group is empty).  we do require
			# that the user has made all the needed groups
			# already, prior to running the script, and make no
			# attempt to make the groups ourselves TODO; we
			# will most likely have switched to qtile before
			# that's needed, this is just a hack for now
			#
			[[ $wm == ratpoison ]] && $rpcmd other

			# since we switched group, the number data for our
			# target window is wrong.  each group has its own
			# numbers.  we must update target info prior to the
			# attempt to renumber, so numbers array update is
			# done correctly for the case where the window
			# already exists (ratpoison will swap them)
			#
			newnum=`$rpcmd "info %n"`
			target_cmp=$saved_grp,$newnum
		fi

		# now we can proceed with the renumber attempt
		$rpcmd "number $saved_num" ||
			bomb "failed to switch number $saved_num"

		# if the destination window already exists,
		# ratpoison will swap it when we do a 'number'
		# so we need to update our gathered title ->
		# num and num -> title maps in-place for the
		# next loop iteration to have correct mapping
		# data after swap
		#
		oldtitle="${windata_number_to_title[$saved_cmp]}"

		if ((debug)); then
			printf "oldtitle: %s\n" "$oldtitle"
			printf "saved_cmp: %s\n" "$saved_cmp"
		fi

		if [[ $oldtitle ]]
		then
			winmap_title_to_number["$saved_title"]=$saved_cmp
			winmap_title_to_number["$oldtitle"]=$target_cmp
			windata_number_to_title[$saved_cmp]="$saved_title"
			windata_number_to_title[$target_cmp]="$oldtitle"
		fi

		if ((debug))
		then
			printf "after-number: %s\n" `print_rpinfo`
			#
			# separate the debug prints with a blank line
			# each loop
			#
			echo
		fi

		sleep 0.4s
	done
}

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

main ()
{
	[[ $1 ]] || bomb "invocation method missing"
	[[ $2 ]] || bomb "browser name missing"
	[[ $2 =~ (chrome|brave|nightly) ]] ||
		bomb "browser name invalid"
	classname1=${classname1[$2]}
	classname2=${classname2[$2]}
	save_file=~/var/$wm/wmpersist-$2.dat
	lock_file=~/var/$wm/$wm.lock

	exec 9>>$lock_file
	flock -w 5 9 || bomb "could not acquire lock"

	# todo: use libsh version of this code
	if function_exists $1
	then $1 "${@:2}"
	else bomb "unimplemented"; fi

	exec 9>&-
}
	
main "$@"