#!/bin/bash
# vi: et sts=4 sw=4 ts=4

USAGE() {
    printf 'Usage: %s [OPTIONS] [--] FILE...\n' \
        "${0##*/}"
}

HELP_MESSAGE() {
    USAGE
    cat <<EOF
Optimize PDF files for size.

  --help, -h        Show this help message.
  --backup=SUFFIX   Keep a backup of un-optimized PDF files.
  --no-backup       Turn off backups (default).
  --careful         Stop if there are warnings in the PDF (default).
  --no-careful      Don't stop if there are warnings in the PDF.
  --force, -f       Force overwriting backups. No effect if --no-backup.
  --no-force        Prompt when overwriting backups (default). No effect if
                      --no-backup.
  --lossy           Let qpdf(1) try to optimize images for size. MAY LEAD TO
                      IMAGE QUALITY DEGRADATION. See qpdf(1) for more
                      information.
  --no-lossy        Don't optimize images (default). This is safer.
  --preserve-timestamp     Copy timestamp from original file.
  --no-preserve-timestamp  Omit timestamp from original file (default).
  --                Terminate options list.

Copyright (C) 2021-2022 Dan Church.
License GPLv3: GNU GPL version 3.0 (https://www.gnu.org/licenses/gpl-3.0.html)
with Commons Clause 1.0 (https://commonsclause.com/).
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.
You may NOT use this software for commercial purposes.
EOF
}

PRESERVE_TIMESTAMP=0
LOSSY=0
KEEP_BACKUP_SUFFIX=
FORCE_OVERWRITE=0
ENCODE_THRU_WARNINGS=0

TEMP_DIR=$(mktemp -d -t "${0##*/}.XXXXXX")
cleanup() {
    rm -fr -- "$TEMP_DIR"
}
trap 'cleanup' EXIT

FILES=()
NO_MORE_FLAGS=0
for ARG; do
    # Assume arguments that don't begin with a - are supposed to be files or other operands
    if [[ $NO_MORE_FLAGS -eq 0 && $ARG = -* ]]; then
        case "$ARG" in
            --backup=*)
                KEEP_BACKUP_SUFFIX=${ARG#*=}
                ;;
            --no-backup)
                KEEP_BACKUP_SUFFIX=
                ;;
            --careful)
                ENCODE_THRU_WARNINGS=0
                ;;
            --no-careful)
                ENCODE_THRU_WARNINGS=1
                ;;
            --force|-f)
                FORCE_OVERWRITE=1
                ;;
            --no-force)
                FORCE_OVERWRITE=0
                ;;
            --lossy)
                LOSSY=1
                ;;
            --no-lossy)
                LOSSY=0
                ;;
            --preserve-timestamp)
                PRESERVE_TIMESTAMP=1
                ;;
            --no-preserve-timestamp)
                PRESERVE_TIMESTAMP=0
                ;;
            --help|-h)
                HELP_MESSAGE
                exit 0
                ;;
            --)
                NO_MORE_FLAGS=1
                ;;
            *)
                printf 'Unrecognized flag: %s\n' \
                    "$ARG" \
                    >&2
                USAGE >&2
                exit 2
                ;;
        esac
    else
        FILES+=("$ARG")
    fi
done

if [[ ${#FILES[@]} -eq 0 ]]; then
    USAGE >&2
    exit 1
fi

file_size() {
    stat \
        --format='%s' \
        --dereference \
        -- \
        "$@" \
        2>/dev/null
}

# produces a human-readable size from the byte count passed to it
hr_size() (
    declare -i BYTES=$1

    #UNITS=(B KB MB GB TB PB EB ZB YB) # shell math can only go so far...
    UNITS=(B KB MB GB TB)
    FACT=1024
    THRESH=9/10
    DECIMALS=1
    DECIMALS_FACTOR=$(( 10 ** DECIMALS ))

    # cycle through units from largest to smallest, exiting when it finds the
    # largest applicable unit
    for (( EXP = ${#UNITS[@]} - 1; EXP > -1; --EXP )); do
        # check if the unit is close enough to the unit's size, within the
        # threshold
        if [[ $BYTES -gt $((FACT ** EXP * $THRESH)) ]]; then
            # we found the applicable unit

            # must multiply by a factor of 10 here to not truncate
            # the given number of decimal places after the point
            HR_VAL=$(( BYTES * DECIMALS_FACTOR / FACT ** EXP ))

            # put the decimal point in
            if [[ $DECIMALS -gt 0 ]]; then
                HR_VAL=$(( HR_VAL / DECIMALS_FACTOR )).$(( HR_VAL % DECIMALS_FACTOR ))
            fi

            HR_UNIT=${UNITS[$EXP]}
            break
        fi
    done

    if [[ -z $HR_UNIT ]]; then
        HR_VAL=$BYTES
        HR_UNIT=${UNITS[0]}
    fi

    printf '%g %s\n' "$HR_VAL" "$HR_UNIT"
)

# copies $2 over to $1 if $2 is smaller than $1
use_smaller() {
    # if `$TEMP' isn't empty and it's of a smaller size than `$FILE',
    # preserve every attribute and replace `$FILE' with `$TEMP'
    local \
        FILE=$1 \
        TEMP=$2 \
        ORIGSIZE \
        TEMPSIZE \
        MV_ARGS=('-v')

    if [[ $FORCE_OVERWRITE -eq 0 ]]; then
        MV_ARGS+=('-i')
    fi

    ORIGSIZE=$(file_size "$FILE")
    TEMPSIZE=$(file_size "$TEMP")

    if [[ -f $TEMP &&
        $TEMPSIZE -gt 0 &&
        $TEMPSIZE -lt $ORIGSIZE
    ]]; then

        # Preserve attributes by copying them from the original file to the
        # temporary one
        chmod \
            --reference="$FILE" \
            -- \
            "$TEMP" &&

        if [[ $PRESERVE_TIMESTAMP -ne 0 ]]; then
            touch \
                --reference="$FILE" \
                -- \
                "$TEMP"
        fi &&

        if [[ $UID -eq 0 ]]; then
            # We are root, so we can chown(1) things
            chown \
                --reference="$FILE" \
                -- \
                "$TEMP"
        fi &&

        if [[ -n $KEEP_BACKUP_SUFFIX ]]; then
            mv "${MV_ARGS[@]}" -- "$FILE" "$FILE$KEEP_BACKUP_SUFFIX"
        fi &&

        cp \
            --preserve=mode,ownership,timestamps \
            -- \
            "$TEMP" \
            "$FILE"

        if [[ $? -ne 0 ]]; then
            printf 'Failed to optimize "%s"!\n' \
                "$FILE" \
                >&2
        fi

    fi

    # Protect against unsuccessful following file writes to our TEMP file
    rm -f -- "$TEMP"

}

QPDF_ARGS=(
    --compression-level=9
    --deterministic-id
    --flatten-rotation
    --object-streams=generate
    --recompress-flate
    --stream-data=compress
)

if [[ $LOSSY -ne 0 ]]; then
    QPDF_ARGS+=(
        --optimize-images
    )
fi
if [[ $ENCODE_THRU_WARNINGS -ne 0 ]]; then
    QPDF_ARGS+=(
        --warning-exit-0
    )
fi

ERRORS=0
FREED_TOTAL=0
for FILE in "${FILES[@]}"; do
    TEMP=$(mktemp -p "$TEMP_DIR" -t 'file.XXXXXX')
    rm -f -- "$TEMP"
    BEGIN_FILESIZE=$(file_size "$FILE")

    if
        qpdf "${QPDF_ARGS[@]}" \
            -- \
            "$FILE" "$TEMP" &&
        use_smaller "$FILE" "$TEMP"; then

        END_FILESIZE=$(file_size "$FILE")
        FREED=$(( BEGIN_FILESIZE - END_FILESIZE ))
        FREED_HR=$(hr_size "$FREED")
        (( FREED_TOTAL += FREED ))

        printf '%s: freed %d bytes (%s)\n' \
            "$FILE" \
            "$FREED" \
            "$FREED_HR"
    else
        (( ++ERRORS ))
    fi
done

FREED_TOTAL_HR=$(hr_size "$FREED_TOTAL")
printf 'all: freed %d bytes (%s)\n' "$FREED_TOTAL" "$FREED_TOTAL_HR"

if [[ $ERRORS -gt 0 ]]; then
    exit 1
fi