#!/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
-r         Reboot now if it's needed.
-h|--help  Print help.

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
}

reboot_not=false
dry_run=false
while true; do
    case $1 in
        -r) reboot_now=true; shift ;;
        -n) dry_run=true; shift ;;
        -h|--help) usage ;;
        *) break ;;
    esac
done

#### 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

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
reboot_script_initialized=false

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

def-e() {
    if $dry_run; then
        e() { echo "+ $@"; }
    else
        e() { echo "+ $@"; "$@"; }
    fi
}

def-e
e swapoff -a

while read devid dev; do
    echo skip=$size
    def-e
    declare -A ptable
    while IFS=: read id start end psize _; do
        [[ $id == [0-9] ]] || continue
        ptable[start$id]=start=${start%%[^0-9]*}
        ptable[end$id]=${end%%[^0-9]*}
        ptable[size$id]=${psize%%[^0-9]*}
    done < <(parted -m /dev/$dev unit MiB print)
    parted /dev/$dev unit MiB print | tee /root/backup_partition_table_$dev
    p() { e parted -a optimal -s -- /dev/$dev unit MiB "$@"; }
    e systemctl stop systemd-cryptsetup@crypt_swap_$dev$swapn
    sleep 1 # dunno if this is needed,
    # but systemd likes to do these kind of things in the background.

    # 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}M /"
    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 echo "$out" | \
            grep "but we have been unable to inform the kernel" &>/dev/null; then
        needs_reboot=true
    fi
    if ! $grow; then
        if $needs_reboot; then
            e systemctl mask dev-mapper-crypt_swap_$dev$swapn.swap
            e systemctl mask systemd-cryptsetup@crypt_swap_$dev$swapn.service
            e() { echo "$@" >> /root/finish-resize; }
            if ! $reboot_script_initialized; then
                reboot_script_initialized=true
                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 systemctl unmask systemd-cryptsetup@crypt_swap_$dev$swapn.service
            e systemctl unmask dev-mapper-crypt_swap_$dev$swapn.swap
        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}M /boot
            e umount /boot
            temp_boot=/root/temp_boot_$dev
            e dd bs=1M if=/dev/$dev$bootn of=$temp_boot \
              count=$((${ptable[size$bootn]} $op_size))
        else
            e umount /boot
        fi
        # if $grow; then
        # ( root ) >< swap >< boot >
        # ( root )( >< swap ) >< boot >
        # ( root )( >< swap )( >< boot )
        # else
        # ( root >< ) swap >< boot >
        # ( root >< )( swap >< ) boot >
        # ( root >< )( swap >< )( 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 mount /boot
            e btrfs fi resize $boot_devid:${op_size}M /boot
        else
            e dd bs=1M if=$temp_boot of=/dev/$dev$bootn
            e mount /boot
        fi
    else
        # if $grow; then  ( root )( >< swap )( boot )
        # else            ( root >< )( swap )( boot )
        pmk $swapn $op_size_rev "" "linux-swap"
        e systemctl start systemd-cryptsetup@crypt_swap_$dev$swapn
    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 "/root/temp_boot_*"
    e rm -f /root/finish-resize
fi

if $needs_reboot; then
    if ! $grow; then
        echo "$0: Reboot, run /root/finish-resize. It's contents:"
        cat /root/finish-resize
    else
        echo "$0: If you want to resize again later, a reboot is required first."
    fi
    if $reboot_now; then
        echo "$0: rebooting now"
        reboot now
        exit
    fi
fi