#!/bin/bash
#
# This is an sdm plugin for: trim-enable
#
# The plugin is called three times: for Phase 0, Phase 1, and post-install.
#

function loadparams() {
    source $SDMPT/etc/sdm/sdm-readparams
}

function istrimenabled() {
    local dev="$1" name aln gran rest
    read name aln gran rest <<< $(lsblk --discard -n $dev)
    [ "$gran" == "0B" ] && return 1 || return 0
}

function trimcfg() {
    local disks="$1" trimenb=0
    # Enable trim on requested disks
    # Based on: https://www.jeffgeerling.com/blog/2020/enabling-trim-on-external-ssd-on-raspberry-pi

    # disks="" [all], "all", or a list of disks ("/dev/sda,/dev/sdb", etc)
    [ "$disks" == "" ] && disks="all"
    while read id diskname
    do
	if [ "$disks" == "all" ] || [[ "$disks" =~ "$diskname" ]]
	then
	    logtoboth "> Plugin $pfx: found disk '$diskname'"
	    id=${id#[}
	    id=${id%]}
	    rawdev=${diskname##/dev/}
	    istrimenabled $diskname && trimenb=0 || trimenb=1
	    pmfile=/sys/block/$rawdev/device/scsi_disk/$id/provisioning_mode 
	    if [ ! -f  $pmfile ]
	    then
		if [ $trimenb -eq 0 ]
		then
		    logtoboth "> Device '$diskname' already has trim enabled"
		else
		    logtoboth "% Device '$diskname' cannot be enabled for trim"
		fi
	    fi
	    if [[ $trimenb -eq 1 ]] && [[ -f $pmfile ]]  # If trim not enabled, but can be
	    then
		lbacount=$(sg_vpd -p bl $diskname 2>/dev/null| grep "unmap LBA" | (read w1 w2 w3 w4 lbc rest; echo $lbc))
		unmapsupp=$(sg_vpd -p lbpv $diskname 2>/dev/null | grep "Unmap command supported" | (read w1 w2 w3 w4 ws ; echo $ws))
		[ "$lbacount" == "" ] && lbacount=0
		[ "$unmapsupp" == "" ] && unmapsupp=0
		if [[ $lbacount -gt 0 ]] && [[ $unmapsupp -eq 1 ]]
		then
		    maxbytes=$((lbacount*512))
		    read pmode < $pmfile
		    if [ "$pmode" != "unmap" ]
		    then
			logtoboth "> Plugin $pfx: Set provisioning mode 'unmap' on '$diskname'"
			csts=$(set -o | grep noclobber | (read vn vs ; echo $vs))
			[ "$csts" == "on" ] && set +o noclobber
			echo "unmap" > $pmfile
			[ "$csts" == "on" ] && set -o noclobber
			logtoboth "> Plugin $pfx: Set discard_max_bytes to $maxbytes"
			echo "$maxbytes" > /sys/block/$rawdev/queue/discard_max_bytes
		    else
			logtoboth "> Plugin $pfx: Device $diskname provisioning mode already set to unmap"
		    fi
		    #
		    # Create udev rule for this disk
		    #
		    udevrule="/etc/udev/rules.d/10-$rawdev-trim.rules"
		    oldudevrule="/etc/udev/rules.d/.10-$rawdev-trim-pre-$(date +'%Y-%m-%d-%H-%M-%S').rules"
		    if [ -f $udevrule ]
		    then
			logtoboth "> Plugin $pfx: Renaming existing udev rule $udevrule to $oldudevrule"
			mv $udevrule $oldudevrule
		    fi
		    idp=""
		    idv=""
		    while read usbdir
		    do
			dskdir="$(find $usbdir/ -name $rawdev -print)"
			if [ "$dskdir" != "" ]
			then
			    #  strip everything after "host" in $dskdir
			    devpath=${dskdir%%/host*}
			    [ ! -L $devpath/port ] && devpath="${devpath%/*}"
			    if [ -L $devpath/port ]
			    then
				idp=$(cat $devpath/idProduct)
				idv=$(cat $devpath/idVendor)
			    else
				logtoboth "? Plugin $pfx: Could not find USB path for $rawdev"
				logtoboth "               Please sudo /usr/local/bin/debug-trim-enable to gather output for maintainer"
				rm -f /usr/local/bin/debug-trim-enable
				trimlog="/root/trim-enable.log"
				cat > /usr/local/bin/debug-trim-enable <<EOF
#!/bin/bash
sudo printf "Debug information for device $diskname\n" > $trimlog
sudo printf "\n** lsusb\n" >> $trimlog
sudo lsusb >> $trimlog
sudo printf "\n** ls -l /etc/udev/rules.d\n" >> $trimlog
sudo ls -l /etc/udev/rules.d >> $trimlog
sudo printf "\n** cat /etc/udev/rules.d/10*trim.rules\n" >> $trimlog
sudo cat /etc/udev/rules.d/10-*-*trim.rules >> $trimlog
sudo printf "\n** sg_vpd -p bl $diskname\n" >> $trimlog
sudo sg_vpd -p bl $diskname  >> $trimlog
sudo printf "\n** sg_vpd -p lbpv $diskname\n" >> $trimlog
sudo sg_vpd -p lbpv $diskname >> $trimlog
echo "Please get the file '$trimlog' to the maintainer"
EOF
				chmod 755 /usr/local/bin/debug-trim-enable
			    fi
			fi
		    done < <(find /sys/bus/usb/devices -maxdepth 1 -type l  -regextype egrep -regex '.*usb[0-9]' -print)
		    if [ "$idv" != "" -a "$idp" != "" ]
		    then
			echo "ACTION==\"add|change\", ATTRS{idVendor}==\"$idv\", ATTRS{idProduct}==\"$idp\", SUBSYSTEM==\"scsi_disk\", ATTR{provisioning_mode}=\"unmap\"" > $udevrule
			echo "ACTION==\"add|change\", KERNEL==\"$rawdev\", SUBSYSTEM==\"block\", RUN+=\"/bin/sh -c 'echo $maxbytes > /sys/block/$rawdev/queue/discard_max_bytes'\"" >> $udevrule
			logtoboth "> Plugin $pfx: Write udev rule $udevrule"
			logtoboth "> Plugin $pfx:  with contents: $(cat $udevrule)"
		    else
			logtoboth "% Plugin $pfx: Did not find idProduct and idVendor for $diskname"
		    fi
		else
		    logtoboth "% Plugin $pfx: Device '$diskname' does not support trim"
		fi
	    fi
	fi
    done < <(lsscsi --brief)
}

# $1 is the phase: "0", "1", or "post-install"
# $2 is the argument list: arg1=val1|arg2=val2|arg3=val3| ...
#
# Main code for the Plugin
#
phase=$1
pfx="$(basename $0)"     #For messages
args="$2"
loadparams

# disks=all, "", or a list of disks (/dev/sda,/dev/sdb, etc)
# To run as boot-time: $1:boot $2:argument list

vldargs="|disks|"
rqdargs=""

if [ "$phase" == "0" ]
then
    #
    # In Phase 0 all references to directories in the image must be preceded by $SDMPT
    #
    logtoboth "* Plugin $pfx: Start Phase 0"
    plugin_getargs $pfx "$args" "$vldargs" "$rqdargs" || exit
    #
    # Print the keys found (example usage). plugin_getargs returns the list of found keys in $foundkeys
    #
    plugin_printkeys
    #plugin_dbgprint "This is how to do a Plugin Debug printout"      # Will only be printed if --plugin-debug specified
    logtoboth "* Plugin $pfx: Complete Phase 0"

elif [ "$phase" == "1" ]
then
    #
    # Phase 1 (in nspawn)
    #
    logtoboth "* Plugin $pfx: Start Phase 1"
    plugin_getargs $pfx "$args" "$vldargs" "$rqdargs"
    #logfreespace "at start of Plugin $pfx Phase 1"
    logtoboth "> Plugin $pfx: Install sg3-utils and lsscsi"
    installpkgsif "sg3-utils lsscsi"
    #logfreespace "at end of $pfx Phase 1"
    logtoboth "* Plugin $pfx: Complete Phase 1"
elif [ "$phase" == "post-install" ]
then
    #
    # Plugin Post-install edits
    #
    logtoboth "* Plugin $pfx: Start Phase post-install"
    plugin_getargs $pfx "$args" "$vldargs" "$rqdargs"
    #logfreespace "at start of Plugin $pfx Phase post-install"
    logtoboth "> Plugin $pfx: Set trim to enable during sdm FirstBoot"
    cat > /etc/sdm/0piboot/020-trim-enable.sh <<EOF
#!/bin/bash
logger "sdm FirstBoot: Enable trim"
$sdmdir/plugins/trim-enable boot "$args"
EOF
    logtoboth "* Plugin $pfx: Complete Phase post-install"
elif [ "$phase" == "boot" ]
then
    # Enable trim on requested disks
    # Based on: https://www.jeffgeerling.com/blog/2020/enabling-trim-on-external-ssd-on-raspberry-pi

    # disks="" [all], "all", or a list of disks ("/dev/sda,/dev/sdb", etc)
    disks="$args"
    [ "$disks" == "" ] && disks="all"
    logtoboth "> Plugin $pfx: Start trim-enable on disks '$disks'"
    trimcfg "$disks"
    #logfreespace "at end of $pfx Custom Phase post-install"
fi
