#!/bin/bash

shopt -s extdebug
bash-trace() {
    # shows function args when: shopt -s extdebug
    local -i argc_index=0 arg frame i start=${1:-1} max_indent=8 indent
    local source
    local extdebug=false
    if [[ $(shopt -p extdebug) == *-s* ]]; then
        extdebug=true
    fi

    for ((frame=0; frame < ${#FUNCNAME[@]}-1; frame++)); do
        argc=${BASH_ARGC[frame]}
        argc_index+=$argc
        ((frame < start)) && continue
        if (( ${#BASH_SOURCE[@]} > 1 )); then
            source="${BASH_SOURCE[frame+1]}:${BASH_LINENO[frame]}:"
        fi
        indent=$((frame-start+1))
        indent=$((indent < max_indent ? indent : max_indent))
        printf "%${indent}s↳%sin \`%s" '' "$source" "${FUNCNAME[frame]}"
        if $extdebug; then
            for ((i=argc_index-1; i >= argc_index-argc; i--)); do
                printf " %s" "${BASH_ARGV[i]}"
            done
        fi
        echo \'
    done
}


errcatch() {
    set -E; shopt -s extdebug
    _err-trap() {
        err=$?
        exec >&2
        echo "${BASH_SOURCE[1]}:${BASH_LINENO[0]}:in \`$BASH_COMMAND' returned $err"
        bash-trace 2
        echo "$0: exiting with code $err"
        exit $err
    }
    trap _err-trap ERR
    set -o pipefail
}

errcatch

[[ $EUID == 0 ]] || sudo "$BASH_SOURCE" "$@"

usage() {
    cat <<EOF
Usage: ${0##*/} [-n] +/-SIZE[g] swap|boot

Assuming Ian Kelling's partition scheme,
Resize swap or boot, expanding or shrinking the root fs and partition to compensate.

-n   Dry run

SIZE is MiB, or if g is specified, GiB.

If using multiple devices, SIZE is applied to each device, so total change is
SIZE * devices.


EOF
    exit $1
}

dry_run=false
case $1 in
    -n) dry_run=true; shift ;;
    -h|--help) usage ;;
esac

#### begin arg error checking ####

if [[ $# != 2 ]]; then
    echo "$0: error: expected 2 arguments"
    usage 1
fi

if [[ $1 != [+-][0-9]* ]]; then
    echo "$0: error: bad 1st arg: $1"
    usage 1
fi

if ! which parted &>/dev/null; then
    echo "$0: error: install parted"
    exit 1
fi

case $2 in swap|boot) : ;; *) echo "$0: error: bad 2nd arg"; usage 1 ;; esac

#### end arg error checking ####


boot=true
[[ $2 == boot ]] || boot=false

#size=${x#[+-]}
op_size=$1 # operator plus size
[[ $op_size != *g ]] || op_size=$(( ${op_size%g} / 1024 ))
size=${op_size#[+-]}

if [[ $op_size == +* ]]; then
    op_size_rev=-$size # rev = reverse
    grow=true
else
    op_size_rev=+$size
    grow=false
fi

##### end command line parsing ########

rootn=1
swapn=2
bootn=3
needs_reboot=false

pmk() {
    part=$1
    start_op=$2
    end_op=$3
    read start end < <(echo ${ptable[$part]})
    p mkpart primary "$4" $((start $start_op)) $((end $end_op))
}


swapoff -a

while read devid dev; do
    if $dry_run; then
        e() { echo "+ $@"; }
    else
        e() { echo "+ $@"; "$@"; }
    fi
    ptable=()
    while IFS=: read id start end _; do
        [[ $id == [0-9] ]] || continue
        end=${end%MiB}
        start=${start%MiB}
        start=${start%.*} # small enough number that parted uses a decimal
        ptable[$id]="$start $end"
    done < <(parted -m /dev/$dev unit MiB print)
    parted /dev/$dev unit MiB print | tee /root/backup_partition
    p() { e parted -a optimal -s -- /dev/$dev unit MiB "$@"; }
    e systemctl stop systemd-cryptsetup@crypt_swap_$dev$swapn
    sleep 1
    # These partition comments seems a little verbose now, but I bet they
    # will be helpfull if I read this in more than a week from now.
    # <> = deleted partition, () = partition
    p rm $swapn # ( root )< swap >( boot )

    root_resize_cmd="e btrfs fi resize $devid:$op_size_rev /"
    if $grow; then $root_resize_cmd; fi
    # if $grow; then
    # < root >< swap >( boot )
    # ( root ) >< swap >( boot )
    # else
    # < root >< swap >( boot )
    # ( root >< ) swap >( boot )

    out=$(p rm $rootn 2>&1)
    echo "$out"

    pmk $rootn "" $op_size_rev
    if ! $grow; then
        if echo "$out" | \
                grep "but we have been unable to inform the kernel" &>/dev/null; then
            e systemctl mask dev-mapper-crypt_swap_$dev$swapn.swap
            e systemctl mask systemd-cryptsetup@crypt_swap_$dev$swapn.service
            if ! $needs_reboot; then
                needs_reboot=true
                echo "$0: reboot and run /root/finish-resize to finish.
The following commands are what will be executed:"
                rm -rf /root/finish-resize
                cat >/root/finish-resize <<'EOF'
#!/bin/bash -x
set -eE -o pipefail
trap 'echo "$0:$LINENO:error: \"$BASH_COMMAND\" returned $?"' ERR
EOF
                chmod +x /root/finish-resize
            fi
            e() { echo "$@" | tee -a /root/finish-resize; }
            e systemctl unmask systemd-cryptsetup@crypt_swap_$dev$swapn.service
            e systemctl unmask dev-mapper-crypt_swap_$dev$swapn.swap
            # todo, disable swap for the next boot. currently have to wait
            # 1:30 for it to fail.
        fi
        $root_resize_cmd
    fi
    if $boot; then
        boot_devid=$(btrfs fi show /boot | \
                            sed -rn "s#^\s*devid\s+(\S+)\s.*$dev[0-9]#\1#p")

        if ! $grow; then
            # shrink boot, move it to a temp file
            e btrfs fi resize $boot_devid:$op_size /boot
            temp_boot=/root/temp_boot_dd
            e dd bs=1M if=/dev/$dev$bootn of=$temp_boot count=$size
        fi
        # if $grow; then
        # ( root ) >< swap >< boot >
        # ( root )( >< swap ) >< boot >
        # ( root )( >< swap )( >< boot )
        # else
        # ( root >< ) swap >< boot >
        # ( root >< )( swap >< ) boot >
        # ( root >< )( swap >< )( boot )
        e umount /boot
        p rm $bootn
        pmk $swapn $op_size_rev $op_size_rev "linux-swap"
        pmk $bootn $op_size_rev ""

        if $grow; then
            e dd bs=1M if=/dev/$dev$bootn of=/dev/$dev$bootn skip=$size
            e btrfs fi resize $boot_devid:$op_size /boot
        else
            e dd bs=1M if=$temp_boot of=/dev/$dev$bootn
        fi
        e mount /boot
    else
        # if $grow; then  ( root )( >< swap )( boot )
        # else            ( root >< )( swap )( boot )
        pmk $swapn $op_size_rev "" "linux-swap"
    fi
done < <(btrfs fi show / | \
                sed -nr 's#^\s*devid\s*(\S+)\s.*_([^_ ]+)[0-9]\s*$#\1 \2#p')

if $boot; then
    e rm -rf $temp_boot
    e rm /root/finish-resize
fi

if $needs_reboot; then
    echo "$0: reminder, reboot then /root/finish-resize"
fi

#for dev in ${devs[@]}; do