#!/bin/sh
set -eu

OUTPUT_DIR="${OUTPUT_DIR:-/scans}"
RESOLUTION="${RESOLUTION:-300}"
MODE="${MODE:-Color}"
SCAN_SOURCE="${SCAN_SOURCE:-Flatbed}"
SCAN_WIDTH="${SCAN_WIDTH:-215.9}"
SCAN_HEIGHT="${SCAN_HEIGHT:-297}"
JPEG_QUALITY="${JPEG_QUALITY:-90}"
SCAN_TIMEOUT="${SCAN_TIMEOUT:-300}"
PUID="${PUID:-0}"
PGID="${PGID:-0}"

mkdir -p "$OUTPUT_DIR"

DEVICE="${SANE_DEVICE:-}"

if [ -z "$DEVICE" ]; then
	DEVICE="$(scanimage -f '%d%n' 2>/dev/null | sed -n '/^hpaio:/{p;q;}')"
fi

if [ -z "$DEVICE" ]; then
	echo "Error: HP scanner not found." >&2
	scanimage -L || true
	exit 2
fi

STAMP="$(date '+%Y-%m-%d_%H-%M-%S')"

TMP_PNM="$(mktemp "/tmp/${STAMP}.pnm.XXXXXX")"
TMP_FIXED="${TMP_PNM}.fixed"
OUTPUT_FILE="${OUTPUT_DIR}/${STAMP}.pdf"
OUTPUT_COMPLETE=0

cleanup() {
	rm -f "$TMP_PNM" "$TMP_FIXED"
	if [ "$OUTPUT_COMPLETE" -ne 1 ]; then
		rm -f "$OUTPUT_FILE"
	fi
}

trap cleanup EXIT HUP INT TERM

SUFFIX=1
while [ -e "$OUTPUT_FILE" ]; do
	OUTPUT_FILE="${OUTPUT_DIR}/${STAMP}_${SUFFIX}.pdf"
	SUFFIX=$((SUFFIX + 1))
done

echo "Scanner: $DEVICE"
echo "Source: $SCAN_SOURCE"
echo "Resolution: $RESOLUTION dpi"
echo "Mode: $MODE"
echo "Area: ${SCAN_WIDTH} x ${SCAN_HEIGHT} mm"
echo "Output: $OUTPUT_FILE"

set +e
timeout "$SCAN_TIMEOUT" scanimage \
	-d "$DEVICE" \
	--source "$SCAN_SOURCE" \
	--resolution "$RESOLUTION" \
	--mode "$MODE" \
	-x "$SCAN_WIDTH" \
	-y "$SCAN_HEIGHT" \
	--format=pnm \
	--output-file="$TMP_PNM"
SCAN_STATUS=$?
set -e

if [ "$SCAN_STATUS" -ne 0 ]; then
	if [ "$SCAN_STATUS" -eq 124 ] || [ "$SCAN_STATUS" -eq 137 ] || [ "$SCAN_STATUS" -eq 143 ]; then
		echo "Error: scanner did not respond within ${SCAN_TIMEOUT} seconds." >&2
	else
		echo "Error: scanimage failed with status ${SCAN_STATUS}." >&2
	fi
	exit "$SCAN_STATUS"
fi

if [ ! -s "$TMP_PNM" ]; then
	echo "Error: scanner returned an empty PNM file." >&2
	exit 3
fi

normalize_pnm_height() {
	MAGIC="$(sed -n '1p' "$TMP_PNM")"
	COMMENT="$(sed -n '2p' "$TMP_PNM")"
	DIMENSIONS="$(sed -n '3p' "$TMP_PNM")"

	if [ "$COMMENT" != '# SANE data follows' ]; then
		echo "Error: scanner returned an unexpected PNM header." >&2
		return 1
	fi

	set -- $DIMENSIONS
	if [ "$#" -ne 2 ]; then
		echo "Error: scanner returned invalid PNM dimensions." >&2
		return 1
	fi

	WIDTH="$1"
	HEIGHT="$2"

	case "${WIDTH}:${HEIGHT}" in
		*[!0-9:]* | :* | *:)
			echo "Error: scanner returned invalid PNM dimensions." >&2
			return 1
			;;
	esac

	case "$MAGIC" in
		P6)
			CHANNELS=3
			HEADER_LINES=4
			;;
		P5)
			CHANNELS=1
			HEADER_LINES=4
			;;
		P4)
			CHANNELS=0
			HEADER_LINES=3
			;;
		*)
			echo "Error: scanner returned unsupported PNM type ${MAGIC}." >&2
			return 1
			;;
	esac

	if [ "$MAGIC" = P4 ]; then
		ROW_BYTES=$(((WIDTH + 7) / 8))
	else
		MAX_VALUE="$(sed -n '4p' "$TMP_PNM")"
		case "$MAX_VALUE" in
			'' | *[!0-9]*)
				echo "Error: scanner returned an invalid PNM color depth." >&2
				return 1
				;;
		esac

		if [ "$MAX_VALUE" -gt 255 ]; then
			SAMPLE_BYTES=2
		else
			SAMPLE_BYTES=1
		fi
		ROW_BYTES=$((WIDTH * CHANNELS * SAMPLE_BYTES))
	fi

	HEADER_BYTES="$(head -n "$HEADER_LINES" "$TMP_PNM" | wc -c)"
	FILE_BYTES="$(wc -c < "$TMP_PNM")"
	PAYLOAD_BYTES=$((FILE_BYTES - HEADER_BYTES))
	EXPECTED_BYTES=$((ROW_BYTES * HEIGHT))

	if [ "$PAYLOAD_BYTES" -eq "$EXPECTED_BYTES" ]; then
		return 0
	fi

	if [ "$PAYLOAD_BYTES" -lt 0 ] || [ $((PAYLOAD_BYTES % ROW_BYTES)) -ne 0 ]; then
		echo "Error: scanner returned a truncated PNM row." >&2
		return 1
	fi

	ACTUAL_HEIGHT=$((PAYLOAD_BYTES / ROW_BYTES))
	MISSING_ROWS=$((HEIGHT - ACTUAL_HEIGHT))

	if [ "$MISSING_ROWS" -ne 1 ]; then
		echo "Error: scanner returned ${ACTUAL_HEIGHT} rows instead of ${HEIGHT}." >&2
		return 1
	fi

	echo "Warning: correcting the HPAIO one-row PNM metadata error." >&2

	{
		head -n 2 "$TMP_PNM"
		printf '%s %s\n' "$WIDTH" "$ACTUAL_HEIGHT"
		if [ "$HEADER_LINES" -eq 4 ]; then
			printf '%s\n' "$MAX_VALUE"
		fi
		tail -c "+$((HEADER_BYTES + 1))" "$TMP_PNM"
	} > "$TMP_FIXED"

	mv "$TMP_FIXED" "$TMP_PNM"
}

normalize_pnm_height

echo "Scan completed."

magick \
	-density "$RESOLUTION" \
	"$TMP_PNM" \
	-compress JPEG \
	-quality "$JPEG_QUALITY" \
	"$OUTPUT_FILE"

if [ ! -s "$OUTPUT_FILE" ]; then
	echo "Error: PDF was not created." >&2
	exit 4
fi

PDF_MAGIC="$(dd if="$OUTPUT_FILE" bs=5 count=1 2>/dev/null)"
if [ "$PDF_MAGIC" != '%PDF-' ]; then
	echo "Error: output file is not a valid PDF." >&2
	exit 5
fi

chown "$PUID:$PGID" "$OUTPUT_FILE" 2>/dev/null || true

OUTPUT_COMPLETE=1
echo "Saved: $OUTPUT_FILE"