#!/usr/bin/env bash
# recompress -- batch convert images to JPEG
#
# Dependencies:
# 	magick, ffmpeg, numfmt, awk, ~/bin/rlisthosts, ~/bin/on
#
# For distributed mode, if SSH multiplexing is enabled then the servers should
# have their 'MaxSessions' increased from the default.

. lib.bash || exit

usage() {
	echo "Usage: $0 <files>"
	echo
	echo_opt "-q <percent>"	"Quality for JPEG compression (default $quality%)"
	echo_opt "-s <percent>|<WxH>" \
				"Scale or dimensions for the output image (default $scale%)"
	echo_opt "-B <color>"	"Background color for transparent PNGs (default $background)"
	echo_opt "-O <arg>"	"Additional option for 'magick convert'"
	echo
	echo_opt "-L <size>"	"Only convert files larger than X megabytes"
	echo_opt "-f"		"Keep converted files even if larger than originals"
	echo_opt "-T"		"Don't try to preserve modification time of the original"
	echo_opt "-d <dir>"	"Put converted files in another directory"
	echo_opt "-x"		"Remove originals after successful conversion"
	echo
	echo_opt "-j <nproc>"	"Number of parallel conversions (default $nproc)"
	echo_opt "-h"		"Distribute conversion across the '@cpu' group"
	echo_opt "-H <hosts>"	"Distribute conversion across specified hosts"
}

# Print out a progress bar, and set the ConEmu taskbar progress.
show_progress() {
	local -i done=$1 total=$2 width=40
	local -i fill=$(( width * done / total ))
	local -i perc=$(( 100 * done / total ))
	local lbar rbar
	printf -v lbar '%*s' $fill ''; lbar=${lbar// /#}
	printf -v rbar '%*s' $(( width-fill )) ''
	printf '%3s%% [%s%s] %s/%s done\r' "$perc" "$lbar" "$rbar" "$done" "$total"
	# Set window title
	printf '\e]0;%s\e\\' "Recompressing $perc% ($done/$total)"
	# Set tmux window name
	if [[ $TERM == @(screen|tmux)* ]]; then
		printf '\ek%s\e\\' "$perc% ($done/$total)"
	fi
	# Set ConEmu progress
	printf '\e]9;4;%d;%d\e\\' 1 "$perc"
}

# Finalize the progress bar, go to next line, and clear the ConEmu progress.
clear_progress() {
	printf '\n'
	printf '\e]0;%s\e\\' ""
	printf '\ek%s\e\\' ""
	printf '\e]9;4;%d;%d\e\\' 0 0
}

dispose() {
	if have trash && [[ -d ~/.local/share/Trash ]]; then
		trash "$@"
	else
		rm -v "$@"
	fi
}

if ! have magick; then
	vdie "magick is not installed"
fi

opt_verbose=0
quality=92
scale=100
keeplarger=0
background='#eeeeee'
minorigsize=0
keepmtime=1
trashorig=0
infiles=()
outfiles=()
infix="conv"
outdir=
declare -A iomap=()
declare -A oimap=()
addargs=()
opt_nproc=0		# -j
opt_hosts=""		# -H

memperjob=100		# wanted MB of free RAM for each worker
hosts=""
nproc=0			# max concurrent workers (total across all hosts)
njobs=0			# active workers
ndone=0			# exited workers (both success and failed)
nfail=0			# failed workers
declare -A hostproc=()	# max workers by host
declare -A hostjobs=()	# active workers by host
declare -A jobfile=()	# input filename by PID
declare -A jobhost=()	# hostname by PID

while getopts ":B:d:fhH:j:L:O:q:s:Txv" OPT; do
	case $OPT in
	B) background="#${OPTARG#'#'}";;
	d) outdir=$OPTARG;;
	f) keeplarger=1;;
	h) opt_hosts="@cpu";;
	H) opt_hosts=$OPTARG;;
	q) quality=${OPTARG%'%'};;
	j) opt_nproc=$OPTARG;;
	L) minorigsize=$OPTARG;;
	O) addargs+=("$OPTARG");;
	s) scale=${OPTARG%'%'};;
	T) keepmtime=0;;
	x) trashorig=1;;
	v) let ++opt_verbose;;
	*) lib:die_getopts;;
	esac
done; shift $((OPTIND-1))

if (( ! $# )); then
	die 0 "no input files specified"
fi

if [[ $outdir ]]; then
	keeplarger=1
fi

if [[ $minorigsize != 0 ]]; then
	minorigsize=$(numfmt --from=iec "${minorigsize%[Bb]}") || exit
	if (( minorigsize < 128 )); then
		warn "assuming $minorigsize is MB rather than bytes"
		minorigsize=$(numfmt --from=iec "${minorigsize}M") || exit
	elif (( minorigsize < 1024 )); then
		warn "assuming $minorigsize is kB rather than bytes"
		minorigsize=$(numfmt --from=iec "${minorigsize}K") || exit
	fi
	info "will skip files smaller than $(numfmt --to=iec "$minorigsize")B"
fi

# Discard inputs which are already converted (to allow 'recompress *.jpg'),
# or those below the specified size limit.
nskipconv=0
nskipsmall=0
nskipother=0
nwarnanim=0
if (( $# > 250 )); then
	# Show signs of life as the 'stat' loop is not instant with 3k inputs
	vmsg "checking $# input files"
fi
for arg in "$@"; do
	oarg=${outdir%/}${outdir:+/}$arg
	if [[ ! -e $arg ]]; then
		err "input file '$arg' does not exist"
	elif [[ ${arg,,} = *.@(zip|webm) ]]; then
		warn "skipping incompatible file '$arg'"
		let ++nskipother
	elif [[ $arg =~ \.conv\.[^.]+$ ]] && [[ $arg -ef $oarg ]]; then
		debug "skipping already converted file '$arg'"
		let ++nskipconv
	elif (( minorigsize && $(stat -c %s "$arg") < minorigsize )); then
		debug "skipping small file '$arg'"
		let ++nskipsmall
	else
		infiles+=("$arg")
	fi
done
if (( errors )); then
	exit 1
fi
if (( nskipconv )); then
	info "$nskipconv arguments ignored as already converted"
fi
if (( nskipsmall )); then
	info "$nskipsmall files skipped as below size threshold"
fi
if (( nskipother )); then
	info "$nskipother files skipped as incompatible"
fi
if (( ! ${#infiles[@]} )); then
	die 0 "no input files"
fi

# Do some additional checks for formats we can't convert
for arg in "${infiles[@]}"; do
	if [[ $arg == *.webp ]]; then
		nframes=$(magick identify -ping "$arg" | wc -l)
		if (( nframes > 1 )); then
			err "file is animated ($nframes frames): $arg"
			let ++nwarnanim
		fi
	fi
done
if (( errors )); then
	exit 1
fi

# Generate imagemagick arguments
addargs+=(-background "$background" -alpha remove)
addargs+=(-quality "$quality")
if (( quality < 20 || quality > 100 )); then
	die "invalid quality '$quality'"
elif (( quality < 90 )); then
	infix="q$quality.$infix"
fi

if [[ $scale == *x* ]]; then
	addargs+=(-resize "$scale")
	infix="resize.$infix"
elif (( scale < 1 || scale > 100 )); then
	die "invalid scale value '$scale'"
elif (( scale < 100 )); then
	addargs+=(-scale "$scale%")
	infix="s$scale.$infix"
fi


# Verify host capacity
if [[ $opt_hosts ]]; then
	hosts=$(rlisthosts "$opt_hosts") || exit
fi
if [[ $hosts ]]; then
	nproc=0
	for hp in $hosts; do
		host=${hp%:*}
		smaxproc=$(ssh "$host" nproc) || exit
		smemfree=$(ssh "$host" \
			"awk '/^MemAvailable:/{print int(\$2/1024)}' /proc/meminfo")
		if [[ $hp == *:* ]]; then
			sproc=${hp#*:}
			if (( sproc < 1 || sproc > smaxproc*2 )); then
				vdie "can't do $sproc jobs on $host ($smaxproc CPUs available)"
			fi
			#if (( opt_nproc )); then
			#	vmsg "host specification '$hp' overrides '-j $opt_nproc'"
			#fi
		elif (( opt_nproc )); then
			sproc=$opt_nproc
			if (( sproc > smaxproc*2 )); then
				vdie "can't do $sproc jobs on $host ($smaxproc CPUs available)"
			fi
		else
			sproc=$smaxproc
			# Add extra jobs to account for slower network I/O
			if [[ $host != "$HOSTNAME" ]]; then
				(( sproc += (sproc+4)/4 ))
			fi
		fi
		if (( smemfree < sproc * memperjob )); then
			vdie "not enough memory for $sproc jobs on $host ($smemfree MB available)"
		fi
		let nproc+=sproc
		hostproc[$host]=$sproc
	done
	vhosts=${!hostproc[*]}
	vmsg "converting ${#infiles[@]} files ($nproc jobs on ${vhosts// /,})"
else
	smaxproc=$(nproc)
	smemfree=$(awk '/^MemAvailable:/{print int($2/1024)}' /proc/meminfo)
	if (( opt_nproc )); then
		sproc=$opt_nproc
		if (( sproc > smaxproc*2 )); then
			vdie "can't do $nproc jobs ($smaxproc CPUs available)"
		fi
	else
		sproc=$smaxproc
	fi
	if (( smemfree < sproc * memperjob )); then
		vdie "not enough memory for $nproc jobs ($smemfree MB available)"
	fi
	nproc=$sproc
	hostproc[$HOSTNAME]=$nproc
	vmsg "converting ${#infiles[@]} files ($nproc local jobs)"
fi

# Decide on output file names upfront (as worker processes cannot easily
# propagate them to the master).
for iname in "${infiles[@]}"; do
	case $iname in
	*.mp3)	ext=mp3;;
	*)	ext=jpg;;
	esac
	oname=${outdir%/}${outdir:+/}${iname%.*}.$infix.$ext
	if [[ ${oimap["$oname"]} && ! ${iomap["$iname"]} ]]; then
		# More than one input (e.g. foo.png & foo.tiff) maps to the
		# same output file, so keep the original extensions.
		xiname=${oimap["$oname"]}
		xoname=${outdir%/}${outdir:+/}$xiname.$infix.$ext
		oname=${outdir%/}${outdir:+/}$iname.$infix.$ext
		iomap["$xiname"]=$xoname
		oimap["$xoname"]=$xiname
		# (Keep the old $oname in the oimap, so that we would detect
		# collisions between 3 items, not only between 2-pairs.)
	fi
	iomap["$iname"]=$oname
	oimap["$oname"]=$iname
done

# Perform the conversion
wait_job() {
	wait -n -p rpid; r=$?
	(( --njobs ))
	(( ++ndone ))
	rhost=${jobhost[$rpid]}; unset jobhost[$rpid]
	rfile=${jobfile[$rpid]}; unset jobfile[$rpid]
	(( --hostjobs[$rhost] ))
	if (( r )); then
		vmsg "job for '$rfile' ($rhost) failed with status $r" >&2
		(( ++nfail ))
	fi
}
trap 'clear_progress' EXIT
show_progress "$ndone" "${#infiles[@]}"
for iname in "${infiles[@]}"; do
	unset host
	for h in "${!hostproc[@]}"; do
		if (( hostjobs[$h] < hostproc[$h] )); then
			host=$h
			break
		fi
	done
	if [[ ! $host ]]; then
		vdie "BUG: did not find a free host (njobs=$njobs, nproc=$nproc)"
	fi

	oname=${iomap["$iname"]}
	case $iname in
	# for libmp3lame, -q:a maps to LAME -V (VBR), -V7 is ~100 kbps
	*.mp3)	cmd=(ffmpeg -loglevel error -i "$iname" -c:a libmp3lame -q:a 7 "$oname");;
	*)	cmd=(magick "$iname" "${addargs[@]}" "$oname");;
	esac

	{
		if [[ $outdir && ! -d ${oname%/*} ]]; then
			mkdir -p "${oname%/*}"
		fi &&
		if [[ $host != "$HOSTNAME" ]]; then
			on "$host" nice "${cmd[@]}"
		else
			nice "${cmd[@]}"
		fi &&
		if (( keepmtime )); then
			touch "$oname" --reference="$iname" 2>/dev/null
		fi; r=$?
		# If any code is added here in the future, always propagate
		# the exit status of the main command.
		exit $r
	} &
	(( ++njobs ))
	jobhost[$!]=$host
	jobfile[$!]=$iname
	(( ++hostjobs[$host] ))
	outfiles+=("$oname")
	show_progress "$ndone" "${#infiles[@]}"
	while (( njobs >= nproc )); do
		wait_job
		show_progress "$ndone" "${#infiles[@]}"
	done
done
while (( njobs > 0 )); do
	wait_job
	show_progress "$ndone" "${#infiles[@]}"
done
clear_progress
trap - EXIT
if (( nfail )); then
	die "$nfail/$ndone jobs returned errors"
fi

# Verify the conversion results
intrash=()
outsmaller=()
inkeep=()
outlarger=()
if (( $# > 250 )); then
	vmsg "verifying output files"
fi
for oname in "${outfiles[@]}"; do
	iname=${oimap["$oname"]}
	if [[ ! -s "$oname" ]]; then
		die "job for '$iname' did not generate output"
	fi
	if (( $(stat -c %s "$oname") >= $(stat -c %s "$iname") - 128*1024 )); then
		inkeep+=("$iname")
		outlarger+=("$oname")
	else
		intrash+=("$iname")
		outsmaller+=("$oname")
	fi
done
if (( ${#outlarger[@]} )) && (( !keeplarger )); then
	warn "${#outlarger[@]} files were larger than originals; deleting."; ((++r))
	rm -f "${outlarger[@]}"
	outfiles=("${outsmaller[@]}")
elif (( ${#outlarger[@]} )); then
	warn "${#outlarger[@]} files are larger than originals"
fi
if (( ! ${#outfiles[@]} )); then
	die "no output files were produced"
fi
final=("${inkeep[@]}" "${outfiles[@]}")

total_size_str() {
	if (( $# )); then
		du -hsc -- "$@" | awk 'END {print $1}'
	else
		echo '0'
	fi
}

avg_size_str() {
	local nbytes=$1 nfiles=$2
	numfmt --to=iec $[nfiles ? (nbytes / nfiles) : 0]
}

orig_total_str=$(total_size_str "${infiles[@]}")
keep_total_str=$(total_size_str "${inkeep[@]}")
conv_total_str=$(total_size_str "${intrash[@]}")
jpeg_total_str=$(total_size_str "${outfiles[@]}")
final_total_str=$(total_size_str "${final[@]}")

orig_total_bytes=$(numfmt --from=iec "$orig_total_str")
keep_total_bytes=$(numfmt --from=iec "$keep_total_str")
conv_total_bytes=$(numfmt --from=iec "$conv_total_str")
jpeg_total_bytes=$(numfmt --from=iec "$jpeg_total_str")
final_total_bytes=$(numfmt --from=iec "$final_total_str")

orig_avg_str=$(avg_size_str $orig_total_bytes ${#infiles[@]})
keep_avg_str=$(avg_size_str $keep_total_bytes ${#inkeep[@]})
conv_avg_str=$(avg_size_str $conv_total_bytes ${#intrash[@]})
jpeg_avg_str=$(avg_size_str $jpeg_total_bytes ${#outfiles[@]})
final_avg_str=$(avg_size_str $final_total_bytes ${#final[@]})

conv_percentage=$[conv_total_bytes ? (100 * jpeg_total_bytes / conv_total_bytes) : 0]
final_percentage=$[orig_total_bytes ? (100 * final_total_bytes / orig_total_bytes) : 0]

echo "Statistics:"
echo "  Orig (all):   $orig_total_str total (${#infiles[@]} files, $orig_avg_str average)"
echo "  Unconverted:  $keep_total_str total (${#inkeep[@]} files, $keep_avg_str average)"
echo "  Orig (conv):  $conv_total_str total (${#intrash[@]} files, $conv_avg_str average)"
echo "  JPEG output:  $jpeg_total_str total (${#outfiles[@]} files, $jpeg_avg_str average)"
echo "  Final result: $final_total_str total (${#final[@]} files, $final_avg_str average)"
echo "  Reduction:    100% -> $conv_percentage% of converted, $final_percentage% of total"

if (( ${#outfiles[@]} > 0 )); then
	i=0
	n=${#outfiles[@]}
	max=15
	echo "Results:"
	for oname in "${outfiles[@]}"; do
		if (( i >= max-1 && n-i > 1 && !opt_verbose )); then
			echo "... [$[n-i] more files, $n total]"
			break
		fi
		let ++i
		if [[ -e "$oname" ]]; then
			osize=$(stat -c %s "$oname")
			iname=${oimap["$oname"]}
			isize=$(stat -c %s "$iname")
			perc=$[100 * osize / isize]
			printf '%6s -> %6s  (%2d%%)  %s\n' \
				"$(numfmt --to=iec $isize)" \
				"$(numfmt --to=iec $osize)" \
				"$perc" \
				"$iname"
		fi
	done
fi

if (( trashorig )); then
	if (( ! ${#intrash[@]} )); then
		echo "No originals to remove (all ${#outlarger[@]} outputs were larger)."
	elif confirm "Remove ${#intrash[@]} originals?"; then
		dispose "${intrash[@]}"
	elif confirm "Remove ${#outsmaller[@]} outputs?"; then
		dispose "${outsmaller[@]}"
	fi
fi

exit $[r > 0]
