#!/usr/bin/bash
#
# The script is part of udev-usb-sync package
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the Affero GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
#    @linux-aarhus - root.nix.dk
#
# configuration : /etc/udev-usb-sync/udev-usb-sync.conf
# triggered by  : /etc/udev/rules.d/99-usb-sync.rules
#
# contributors: @megavolt (Manjaro Forum)
#             : @linux-aarhus (Manjaro Forum)
# inspired by : @kwg (EndeavourOS Forum)
#
# Arguments provided by udev rule
#     $1: usb block device
#     $2: usb bandwidth reported by device
#

VERSION='1.1'

START=$(date +%s%N)

if [[ -z $1 || -z $2 ]]; then
    echo ":: Script: udev-usb-sync version $VERSION"
    echo ":: See <https://gitlab.manjaro.org/applications/udev-usb-sync>"
    exit
fi

set -euo pipefail

LANG=C
LC_NUMERIC=C

AUTOCALC=${AUTOCALC:-1}
CONFIG='/etc/udev-usb-sync/udev-usb-sync.conf'
[[ -f $CONFIG ]] && source $CONFIG

LOG=${LOG:-1}
LOGFILE="/var/log/udev-usb-sync.log"

BLOCKDEVICE="$1"
SPEED="$2"

if [[ $LOG == 1 ]]; then
    log() { printf '%(%F %T)T [%s] %s\n' -1 "$1" "${*:2}"; }
    exec > >(tee -a $LOGFILE) 2>&1
fi

log INFO "$(uname -r)"

if [[ -f "$CONFIG" ]]; then
    log INFO "Configuration: $CONFIG"
else
    log WARNING "No configuration file found."
fi

if [[ $AUTOCALC == 1 ]]; then 
    log INFO "Autocalculation is on."
else
    log WARNING "Autocalculation is off. Fallback to strict 16MB buffer size."
fi

log INFO "Called with -> $BLOCKDEVICE $SPEED"

# disable write cache for device if possible

if command -v hdparm &>/dev/null; then
    hdparm -W 0 "/dev/$BLOCKDEVICE" &>/dev/null && log INFO "Cache disabled." || log WARNING "hdparm failed."
else
    log INFO "hdparm not installed, skipping."
fi

# the following rules is introduced with kernel 6.2
# https://docs.kernel.org/admin-guide/abi-testing.html#abi-sys-class-bdi-bdi-max-bytes

bdi_state() { grep . /sys/block/$BLOCKDEVICE/bdi/{max_bytes,max_ratio,strict_limit} | sed 's|.*/||' | paste -sd ';'; }
b2mb() { printf '%.0fMB' $(bc -l <<< "(${1} / 1024 / 1024 )"); }

log BDI "$(bdi_state)"

case $AUTOCALC in 
    0)  # apply 16M as max_bytes
        echo 16777216 > /sys/block/$BLOCKDEVICE/bdi/max_bytes
    ;;
    1) BUFFER_TIME=${BUFFER_TIME:-"0.05"}
       SAFETY_FACTOR=${SAFETY_FACTOR:-"1.3"}
       BUFFER_SIZE=$(printf '%.0f' "$(bc -l <<< "(($SPEED / 8) * $BUFFER_TIME * $SAFETY_FACTOR) * 1024 * 1024")")
       # for x in 12 480 5000 10000; do echo -n "$x -> " ;printf "%.0f\n" ` echo "(($x / 8) * 0.05 * 1.3) * 1024 * 1024" | bc`; done
       # 62915
       # 4089446
       # 42593157
       # 85196800
       # apply calculated buffer size
       OLD_BUFFER=$(cat /sys/block/$BLOCKDEVICE/bdi/max_bytes)
       log BUFFER "old -> ${OLD_BUFFER} bytes [$(b2mb ${OLD_BUFFER})]"
       log BUFFER "cal -> ${BUFFER_SIZE} bytes [$(b2mb ${BUFFER_SIZE})]"
       echo "$BUFFER_SIZE" > /sys/block/$BLOCKDEVICE/bdi/max_bytes
       NEW_BUFFER=$(cat /sys/block/$BLOCKDEVICE/bdi/max_bytes)
       log BUFFER "new -> ${NEW_BUFFER} bytes [$(b2mb ${NEW_BUFFER})]"
    ;;
esac

# https://docs.kernel.org/admin-guide/abi-testing.html#abi-sys-class-bdi-bdi-strict-limit
# apply strict limit
if echo 1 > /sys/block/$BLOCKDEVICE/bdi/strict_limit; then
    log INFO "Strict limit applied."
fi

log BDI "$(bdi_state)"

END=$(date +%s%N)
DURATION_MS=$(( (END - START) / 1000000 ))
DURATION_SC=$(printf '%d.%03d' $(( DURATION_MS / 1000 )) $(( DURATION_MS % 1000 )))
log INFO "Duration: ${DURATION_MS}ms (${DURATION_SC}s)"
