#!/bin/bash
# adapted for wayland (grim/slurp/wl-copy) from
# https://github.com/tatsumoto-ren/dotfiles/blob/main/.local/bin/maimocr

set -e
[ -n "$DEBUG" ] && set -x

TESSDATA_PREFIX=$HOME/.local/share/tessdata

langs=
method=tesseract
while getopts l:m:h opt; do
	case "$opt" in
	l) langs="${OPTARG}" ;;
	m) method="${OPTARG}" ;;
	h | ?)
		printf "usage: %s [-h]\n" "$(basename "$0")"
		printf "\n"
		printf "  -l  LANG[+LANG]\n"
		printf "  -m  METHOD\n"
		printf "\n"
		exit 2
		;;
	esac
done
shift $((OPTIND - 1))

for c in curl unzip grim slurp convert wl-copy; do
	if ! command -v "$c" >/dev/null; then
		notify "$c must be installed."
		exit 1
	fi
done

notify() {
	echo "$*"
	notify-send "㊗️ grimOCR" "$*"
}

notify_ask() {
	notif_args="gdbus call --session \
		--dest org.freedesktop.Notifications \
		--object-path /org/freedesktop/Notifications"
	ok_action_key=ok__$(mktemp -u | cut -d. -f2)
	cancel_action_key=cancel__$(mktemp -u | cut -d. -f2)

	APP_NAME=grimocr
	REPLACE_ID=0
	ICON=
	SUMMARY=grimOCR
	BODY="$*"
	ACTIONS="[\"$ok_action_key\", \"OK\", \"$cancel_action_key\", \"Cancel\"]"
	HINTS='{}'
	EXPIRE_TIME=0

	id=$($notif_args --method org.freedesktop.Notifications.Notify -- \
		"$APP_NAME" "$REPLACE_ID" "$ICON" "$SUMMARY" "$BODY" \
		"$ACTIONS" "$HINTS" "int32 $EXPIRE_TIME" | sed 's/(uint32 \([0-9]\+\),)/\1/g')

	pidfile=$(mktemp /tmp/grimocr_notification_monitor.XXXXX)
	result=$(mktemp /tmp/grimocr_notification_result.XXXXX)
	trap '{ kill "$(cat "$pidfile")" ; rm -- "$pidfile" "$result"; }' EXIT
	(
		gdbus monitor --session \
			--dest org.freedesktop.Notifications \
			--object-path /org/freedesktop/Notifications &
		printf "%d" $! >&3
	) 3>"$pidfile" |
		while read -r line; do
			closed_id=$(printf "%s" "$line" | sed '/^\/org\/freedesktop\/Notifications: org.freedesktop.Notifications.NotificationClosed (uint32 \([0-9]\+\), uint32 [0-9]\+)$/!d;s//\1/')
			if [ -n "$closed_id" ] && [ "$closed_id" = "$id" ]; then
				$notif_args --method org.freedesktop.Notifications.CloseNotification "$id" >/dev/null
				break
			fi
			action_invoked=$(printf "%s" "$line" | sed "/\/org\/freedesktop\/Notifications: org.freedesktop.Notifications.ActionInvoked (uint32 \([0-9]\+\), '\(.*\)')\$/!d;s//\1:\2/")
			if [ -n "$action_invoked" ]; then
				invoked_id=$(printf "%s" "$action_invoked" | cut -d: -f1)
				action_id=$(printf "%s" "$action_invoked" | cut -d: -f2)
				if [ "$invoked_id" = "$id" ]; then
					case "$action_id" in
					"$ok_action_key") printf "0" >"$result" ;;
					"$cancel_action_key") printf "1" >"$result" ;;
					esac
					$notif_args --method org.freedesktop.Notifications.CloseNotification "$id" >/dev/null
					break
				fi
			fi
		done
	return "$(cat "$result")"
}

download_tessdata() {
	osd_url="https://github.com/tesseract-ocr/tessdata_best/blob/main/osd.traineddata"
	eng_url="https://github.com/tesseract-ocr/tessdata_best/raw/main/eng.traineddata"
	jpn_url="https://g33k.se/_matrix/media/r0/download/g33k.se/chvDlHzjiXgDEdeGTFnyOExv"

	if ! notify_ask "Downloading Tesseract data?"; then
		notify "Download aborted."
		exit 0
	else
		notify "Download started."
	fi

	# download osd tessdata
	curl -LsSf -o "$TESSDATA_PREFIX/osd.tessdata" "$osd_url"
	curl -LsSf -o "$TESSDATA_PREFIX/eng.tessdata" "$eng_url"

	# download jpn tessdata
	jpn_zip_path=$(mktemp /tmp/grimocr_tesseract_data.XXXXX)
	trap '{ rm -- "$jpn_zip_path"; }' EXIT
	curl -LsSf -o "$jpn_zip_path" "$jpn_url"
	mkdir -p -- "$TESSDATA_PREFIX"
	unzip -o -j "$jpn_zip_path" -d "$TESSDATA_PREFIX"

	notify "Download complete. Run again..."
}

take_screenshot() {
	# https://tesseract-ocr.github.io/tessdoc/ImproveQuality.html#missing-borders
	grim -g "$(slurp)" -t png -l 1 - |
		convert png:- -alpha off -bordercolor White -border 10x10 png:-

	notify "recognizing..." >/dev/null
}

installed_languages() {
	if [ -n "$langs" ]; then
		printf "%s" "$langs"
	else
		find "$TESSDATA_PREFIX" -type f -name '*.traineddata' -printf '%f+' |
			sed 's|\.traineddata||g; s|\+$||'
	fi
}

tesseract_recognize() {
	# https://tesseract-ocr.github.io/tessdoc/Command-Line-Usage.html
	tesseract stdin stdout \
		--tessdata-dir "$TESSDATA_PREFIX" \
		-l "$(installed_languages)" \
		--psm 1 |
		tr -d '\f'

	# NOTE: by removing this, we can make it work for other langs too
	# tr -d ' ' |
	# tr -d '\n' |
	# grep -P '[ｦ-ﾟァ-ヶぁ-ゞＡ-ｚ０-９ｧ-ﾝﾞﾟァ-ンぁ-ん一-龯]+'
}

mangaocr_recognize() {
	python3 -c "$(
		cat <<-EOF
			from manga_ocr import MangaOcr

			text = MangaOcr()("/dev/stdin")
			print(text)
		EOF
	)"
}

recognize() {
	case "${method}" in
	tesseract) tesseract_recognize ;;
	mangaocr) mangaocr_recognize ;;
	esac
}

run_ocr() {
	result=$(take_screenshot | recognize 2>/dev/null)
	if [ $? -eq 0 ]; then
		printf "%s" "$result" | wl-copy -n
		notify "$result"
	else
		notify "Failed."
	fi
}

## main

case "${method}" in
tesseract)
	if ! command -v tesseract >/dev/null; then
		notify "tesseract must be installed."
		exit 1
	fi
	if [ ! -d "$TESSDATA_PREFIX" ]; then
		download_tessdata
	fi
	;;
mangaocr)
	if ! command -v python3 >/dev/null; then
		notify "python3 must be installed."
		exit 1
	fi
	notify "setting up manga-ocr..."
	script_path="$(dirname "$(readlink -f "$0")")"
	venv_path="${script_path}"/"$(basename "$0")".mangaocr
	python3 -m venv "${venv_path}"
	source "${venv_path}"/bin/activate
	pip3 install -r "${script_path}"/"$(basename "$0")".requirements >/dev/null 2>&1
	;;
*)
	notify "unknown METHOD ${method}"
	exit 1
	;;
esac

notify "select some japanese text"
run_ocr