#!/bin/sh
# tlp-func-disk - Storage Device and Filesystem Functions
#
# Copyright (c) 2020 Thomas Koch <linrunner at gmx.net> and others.
# This software is licensed under the GPL v2 or later.

# Needs: tlp-func-base

# shellcheck disable=SC2086

# ----------------------------------------------------------------------------
# Constants

readonly AHCID=$PCID'/*/ata*'
readonly BLOCKD='/sys/block/nvme* /sys/block/sd*'

readonly DISK_TYPES_NO_APM_CHANGE="usb ieee1394"
readonly DISK_NOP_WORDS="_ keep"

# ----------------------------------------------------------------------------
# Functions

# --- Device Helpers

get_disk_dev () { # translate disk id to device (sdX)
    # $1: id or dev;
    # retval: $_disk_dev, $_disk_id, $_disk_type
    local bus dpath path

    if [ -h /dev/disk/by-id/$1 ]; then
        # $1 is disk id
        _disk_id=$1
        _disk_dev=$(printf '%s' "$_disk_id" | sed -r 's/-part[1-9][0-9]*$//')
        _disk_dev=$(readlink /dev/disk/by-id/$_disk_dev)
        _disk_dev=${_disk_dev##*/}
    else
        # $1 is disk dev
        _disk_dev=$1
        _disk_id=""
    fi

    # determine device type (bus)
    if [ -b /dev/$_disk_dev ]; then
        path="$($UDEVADM info -q property /dev/$_disk_dev 2>/dev/null | sed -n 's/^ID_PATH=//p')"
        bus="$($UDEVADM  info -q property /dev/$_disk_dev 2>/dev/null | sed -n 's/^ID_BUS=//p')"
        dpath=
        case "$path" in
            pci-*-nvme-*)     _disk_type="nvme" ;;
            pci-*-ata-*)      _disk_type="ata"  ;;
            pci-*-usb-*)      _disk_type="usb"  ;;
            pci-*-ieee1394-*) _disk_type="ieee1394" ;;
            *) case "$bus" in
                nvme)      _disk_type="nvme" ;;
                ata)       _disk_type="ata"  ;;
                usb)       _disk_type="usb"  ;;
                ieee1394)  _disk_type="ieee1394" ;;

                *)
                    dpath="$($UDEVADM info -q property /dev/$_disk_dev 2>/dev/null | sed -n 's/^DEVPATH=//p')"
                    dpath=${dpath##*/}
                    case "$dpath" in
                        nvme*) _disk_type="nvme" ;;
                        *)     _disk_type="unknown" ;;
                    esac
                    ;;
            esac
        esac
        echo_debug "disk" "get_disk_dev($1): dev=$_disk_dev; type=$_disk_type; path=$path; bus=$bus; dpath=$dpath"
        return 0
    else
        _disk_type="none"
        echo_debug "disk" "get_disk_dev($1).missing"
        return 1
    fi
}

check_disk_apm () { # check if disk device supports advanced pm
    # $1: dev; rc: 0=yes/1=no

    if $HDPARM -i $1 2> /dev/null | grep -q 'AdvancedPM=yes'; then
        return 0
    else
        return 1
    fi
}

show_disk_ids () { # show disk id's
    local dev

    { # iterate SATA and NVMe disks
      # shellcheck disable=SC2010
        for dev in $(ls /dev/disk/by-id/ | grep -E '^(ata|ieee1394|nvme|usb)' | grep -E -v '(^nvme-eui|\-part[1-9]+)'); do
            if [ -n "$dev" ]; then
                get_disk_dev $dev && printf '%s: %s\n' "$_disk_dev" "$_disk_id"
            fi
        done
    } | sort

    return 0
}

# --- Disk APM Features

set_disk_apm_level () { # set disk apm level
    # $1: 0=ac mode, 1=battery mode

    local pwrmode="$1"
    local dev log_message

    # quit when disabled
    if [ -z "$DISK_DEVICES" ]; then
        echo_debug "disk" "set_disk_apm_level($pwrmode).disabled"
        return 0
    fi

    # set @argv := apmlist (blanks removed - relying on a sane $IFS)
    if [ "$pwrmode" = "1" ]; then
        set -- $DISK_APM_LEVEL_ON_BAT
    else
        set -- $DISK_APM_LEVEL_ON_AC
    fi

    # quit if empty apmlist
    if [ $# -eq 0 ]; then
        echo_debug "disk" "set_disk_apm_level($pwrmode).not_configured"
        return 0
    fi

    echo_debug "disk" "set_disk_apm_level($pwrmode)"

    # pairwise iteration DISK_DEVICES[1,n], apmlist[1,m]; m > 0
    #  for j in [1,n]: disk_dev[j], apmlist[min(j,m)]
    #
    for dev in $DISK_DEVICES; do
        : ${1:?BUG: broken DISK_APM_LEVEL list handling}

        if get_disk_dev $dev; then
            log_message="set_disk_apm_level($pwrmode): $_disk_dev [$_disk_id] $1"

            if ! check_disk_apm /dev/$_disk_dev || wordinlist "$_disk_type" "$DISK_TYPES_NO_APM_CHANGE"; then
                echo_debug "disk" "${log_message}; not supported"
            elif wordinlist "$1" "$DISK_NOP_WORDS"; then
                echo_debug "disk" "${log_message}; keep as is"
            else
                $HDPARM -B $1 /dev/$_disk_dev > /dev/null 2>&1
                echo_debug "disk" "${log_message}; rc=$?"
            fi
        fi
        # last entry in apmlist applies to all remaining disks
        [ $# -lt 2 ] || shift
    done

    return 0
}

set_disk_spindown_timeout () { # set disk spindown timeout
    # $1: 0=ac mode, 1=battery mode

    local pwrmode="$1"
    local dev log_message

    # quit when disabled
    if [ -z "$DISK_DEVICES" ]; then
        echo_debug "disk" "set_disk_spindown_timeout($pwrmode).disabled"
        return 0
    fi

    # set @argv := timeoutlist
    if [ "$pwrmode" = "1" ]; then
        set -- $DISK_SPINDOWN_TIMEOUT_ON_BAT
    else
        set -- $DISK_SPINDOWN_TIMEOUT_ON_AC
    fi

    # quit if empty timeoutlist
    if [ $# -eq 0 ]; then
        echo_debug "disk" "set_disk_spindown_timeout($pwrmode).not_configured"
        return 0
    fi

    echo_debug "disk" "set_disk_spindown_timeout($pwrmode)"

    # pairwise iteration DISK_DEVICES[1,n], timeoutlist[1,m]; m > 0
    #  for j in [1,n]: disk_dev[j], timeoutlist[min(j,m)]
    #
    for dev in $DISK_DEVICES; do
        : ${1:?BUG: broken DISK_SPINDOWN_TIMEOUT list handling}

        if get_disk_dev $dev; then
            log_message="set_disk_spindown_timeout($pwrmode): $_disk_dev [$_disk_id] $1"

            if wordinlist "$1" "$DISK_NOP_WORDS"; then
                echo_debug "disk" "${log_message}; keep as is"
            else
                $HDPARM -S $1 /dev/$_disk_dev > /dev/null 2>&1
                echo_debug "disk" "${log_message}; rc=$?"
            fi
        fi
        # last entry in timeoutlist applies to all remaining disks
        [ $# -lt 2 ] || shift
    done

    return 0
}

spindown_disk () { # stop spindle motor -- $1: dev
    $HDPARM -y /dev/$1 > /dev/null 2>&1

    return 0
}

set_disk_iosched () { # set disk io scheduler
    local dev log_message

    # quit when disabled
    if [ -z "$DISK_DEVICES" ]; then
        echo_debug "disk" "set_disk_iosched.disabled"
        return 0
    fi

    # set @argv := schedlist
    set -- $DISK_IOSCHED

    # quit if empty schedlist
    if [ $# -eq 0 ]; then
        echo_debug "disk" "set_disk_iosched.not_configured"
        return 0
    fi

    echo_debug "disk" "set_disk_iosched"

    # pairwise iteration DISK_DEVICES[1,n], schedlist[1,m]; m > 0
    #  for j in [1,min(n,m)]   : disk_dev[j], schedlistj]
    #  for j in [min(n,m)+1,n] : disk_dev[j], %keep
    for dev in $DISK_DEVICES; do
        local sched schedctrl
        if get_disk_dev $dev; then
            # get sched from argv, use "keep" when list is too short
            sched=${1:-keep}
            schedctrl="/sys/block/$_disk_dev/queue/scheduler"
            log_message="set_disk_iosched: $_disk_dev [$_disk_id] $sched"

            if [ ! -f $schedctrl ]; then
                echo_debug "disk" "${log_message}; not supported"
            elif wordinlist "$sched" "$DISK_NOP_WORDS"; then
                echo_debug "disk" "${log_message}; keep as is"
            else
                write_sysf "$sched" $schedctrl
                echo_debug "disk" "${log_message}; rc=$?"
            fi
        fi
        # using "keep" when argv is empty
        [ $# -eq 0 ] || shift
    done

    return 0
}

# --- Power Saving

set_sata_link_power () { # set ahci link power management
    # $1: 0=ac mode, 1=battery mode

    local pm="$1"
    local host host_bl hostid linkpol pwr rc
    local pwrlist=""
    local ctrl_avail="0"

    if [ "$pm" = "1" ]; then
        pwrlist=${SATA_LINKPWR_ON_BAT:-}
    else
        pwrlist=${SATA_LINKPWR_ON_AC:-}
    fi

    if [ -z "$pwrlist" ]; then
        # do nothing if unconfigured
        echo_debug "pm" "set_sata_link_power($pm).not_configured"
        return 0
    fi

    # ALPM blacklist
    host_bl=${SATA_LINKPWR_BLACKLIST:-}

    # copy configured values to args
    set -- $pwrlist
    # iterate SATA hosts
    for host in /sys/class/scsi_host/host* ; do
        linkpol=$host/link_power_management_policy
        if [ -f $linkpol ]; then
            hostid=${host##*/}
            if wordinlist "$hostid" "$host_bl"; then
                # host blacklisted --> skip
                echo_debug "pm" "set_sata_link_power($pm).black: $host"
                ctrl_avail="1"
            else
                # host not blacklisted --> iterate all configured values
                for pwr in "$@"; do
                    write_sysf "$pwr" $linkpol; rc=$?
                    echo_debug "pm" "set_sata_link_power($pm).$pwr: $host; rc=$rc"
                    if [ $rc -eq 0 ]; then
                        # write successful --> goto next host
                        ctrl_avail="1"
                        break
                    else
                        # write failed --> don't use this value for remaining hosts
                        # and try next value
                        shift
                    fi
                done
            fi
        fi
    done

    [ "$ctrl_avail" = "0" ] && echo_debug "pm" "set_sata_link_power($pm).not_available"
    return 0
}

set_ahci_runtime_pm () { # set ahci runtime power management
    # $1: 0=ac mode, 1=battery mode

    local control device ec timeout rc

    if [ "$1" = "1" ]; then
        control=${AHCI_RUNTIME_PM_ON_BAT:-}
    else
        control=${AHCI_RUNTIME_PM_ON_AC:-}
    fi

    # calc timeout in millisecs
    timeout="$AHCI_RUNTIME_PM_TIMEOUT"
    [ -z "$timeout" ] || timeout=$((timeout * 1000))

    # check values
    case "$control" in
        on|auto)       ;;
        *) control="" ;; # invalid input --> unconfigured
    esac

    if [ -z "$control" ]; then
        # do nothing if unconfigured
        echo_debug "pm" "set_ahci_runtime_pm($1).not_configured"
        return 0
    fi

    # when timeout is unconfigured we're done here
    if [ -z "$timeout" ]; then
        echo_debug "pm" "set_ahci_runtime_pm($1).timeout_not_configured"
        return 0
    fi

    # iterate disks
    ec=0
    for device in $BLOCKD; do
        if [ -f ${device}/device/power/control ]; then
            # write timeout first because writing "auto" with the default
            # timeout -1 still active will lockup the machine!
            rc=0
            if write_sysf "$timeout" ${device}/device/power/autosuspend_delay_ms; then
                # writing timeout was successful --> proceed with activation;
                if ! write_sysf "$control" ${device}/device/power/control; then
                    # writing control failed
                    rc=2
                    ec=$((ec+1))
                fi
            else
                # writing timeout failed (or file not present)
                rc=1
                ec=$((ec+1))
            fi
            echo_debug "pm" "set_ahci_runtime_pm($1).$control: disk=$device timeout=$timeout; rc=$rc"
        fi
    done

    if [ $ec -gt 0 ]; then
        # quit: do not touch hosts if write failed for at least one disk
        echo_debug "pm" "set_ahci_runtime_pm($1).quit; ec=$ec"
        return 0
    fi

    # iterate hosts
    for device in $AHCID; do
        write_sysf "$control" ${device}/power/control
        echo_debug "pm" "set_ahci_runtime_pm($1).$control: host=$device; rc=$?"
    done

    return 0
}

# --- Filesystem Parameters

set_laptopmode () { # set kernel laptop mode -- $1: 0=ac mode, 1=battery mode
    check_sysfs "set_laptopmode" "/proc/sys/vm/laptop_mode"

    local isec

    if [ "$1" = "1" ]; then
        isec="$DISK_IDLE_SECS_ON_BAT"
    else
        isec="$DISK_IDLE_SECS_ON_AC"
    fi
    # replace with empty string if non-numeric chars are contained
    isec=$(printf '%s' "$isec" | grep -E '^[0-9]+$')

    if [ -z "$isec" ]; then
        # do nothing if unconfigured or non numeric value
        echo_debug "pm" "set_laptopmode($1).not_configured"
        return 0
    fi

    write_sysf "$isec" /proc/sys/vm/laptop_mode
    echo_debug "pm" "set_laptopmode($1): $isec; rc=$?"

    return 0
}

set_dirty_parms () { # set filesystem buffer params
    # $1: 0=ac mode, 1=battery mode
    # concept from laptop-mode-tools

    local age cage df ec

    check_sysfs "set_dirty_parms" "/proc/sys/vm"

    if [ "$1" = "1" ]; then
        age=${MAX_LOST_WORK_SECS_ON_BAT:-0}
    else
        age=${MAX_LOST_WORK_SECS_ON_AC:-0}
    fi

    # calc age in centisecs, non numeric values result in "0"
    cage=$((age * 100))

    if [ "$cage" = "0" ]; then
        # do nothing if unconfigured or invalid age
        echo_debug "pm" "set_dirty_parms($1).not_configured"
        return 0
    fi

    ec=0
    for df in /proc/sys/vm/dirty_writeback_centisecs \
             /proc/sys/vm/dirty_expire_centisecs \
             /proc/sys/fs/xfs/age_buffer_centisecs \
             /proc/sys/fs/xfs/xfssyncd_centisecs; do
        if [ -f "$df" ] && ! write_sysf "$cage" $df; then
            echo_debug "pm" "set_dirty_parms($1).write_error: $df $cage; rc=$?"
            ec=$((ec+1))
        fi
    done
    # shellcheck disable=SC2043
    for df in /proc/sys/fs/xfs/xfsbufd_centisecs; do
        if [ -f "$df" ] && ! write_sysf "3000" $df; then
            echo_debug "pm" "set_dirty_parms($1).write_error: $df 3000; rc=$?"
            ec=$((ec+1))
        fi
    done
    echo_debug "pm" "set_dirty_parms($1): $cage; ec=$ec"

    return 0
}
