#!/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 ]] || exec sudo "$BASH_SOURCE" "$@"

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

Assuming Ian Kelling's partition scheme, Resize swap or boot, expanding
or shrinking the root fs and partition to compensate. If it changes
the partition tables incorrectly, the originals are stored in
/root/backup_partition_table_<device_names>.

-n         Dry run
-r         Reboot now if it's needed.
-h|--help  Print help and exit.

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.

Note: Uses GNU getopt options parsing style
EOF
    exit $1
}

reboot_not=false
dry_run=false

temp=$(getopt -l help rnh "$@") || usage 1
eval set -- "$temp"
while true; do
    case $1 in
        -r) reboot_now=true; shift ;;
        -n) dry_run=true; shift ;;
        -h|--help) usage ;;
        --) shift; break ;;
        *) echo "$0: Internal error!" ; exit 1 ;;
    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
if [[ $op_size == *g ]]; then
    op_size=${op_size%g}
    size=${op_size#?}
    op=${op_size%size}
    size=$(( ${size} * 1024 ))
    op_size=$op$size
else
    size=${op_size#?}
fi

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() { # partition make
    part=$1
    start_op=$2
    end_op=$3
    fs_type="$4"

    # This fails outside a vm, but actually succeeds. also prints this
    # message:
    # Error: Partition(s) 2 on /dev/sda have been written, but
    # we have been unable to inform the kernel of the change, probably
    # because it/they are in use.  As a result, the old partition(s)
    # will remain in use.  You should reboot now before making further
    # changes.

    p mkpart primary "$fs_type" \
      $((${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
    if [[ $dev != /dev/mapper/* ]]; then
        # older oses, it points to /dev/dm-x
        dev=$(dmsetup info $dev | sed -rn 's/^\s*Name:\s*(\S*)/\1/p')
    else
        dev=${dev#/dev/mapper/}
    fi
    dev=${dev#crypt_dev_}
    r='-part[0-9]*$'
    [[ $dev =~ $r ]] ||:
    dev=${dev%$BASH_REMATCH}
    devpath=/dev/disk/by-id/$dev
    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 $devpath unit MiB print)
    parted $devpath unit MiB print | tee /root/backup_partition_table_$dev
    p() { e parted -a optimal -s -- $devpath unit MiB "$@"; }
    unit=systemd-cryptsetup@crypt_dev_$dev-part$swapn
    # note systemctl show can test if a unit exists.
    e systemctl stop $unit || [[ $? == 5 ]] # fail due to not loaded
    # there is a bug in jessie. this and the .swap unit are
    # generated from /etc/fstab, and it escapes - to x2d, then doesn't escape it
    # when looking for the file to use as swap. so, no swap is working on jessie
    # right now.
    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 $?" >&2' 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
        # non by-id path, to match what btrfs fi show will tell us
        boot_dev_path=$(readlink -f $devpath-part$bootn)
        boot_devid=$(btrfs fi show /boot | \
                         sed -rn "s#^\s*devid\s+(\S+)\s.*$boot_dev_path#\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=$boot_dev_path 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=$boot_dev_path of=$boot_dev_path skip=$size
            e mount /boot
            e btrfs fi resize $boot_devid:${op_size}M /boot
        else
            e dd bs=1M if=$temp_boot of=$boot_dev_path
            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+).* path (.*)$#\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