#!/bin/bash
# -*- coding: utf-8; indent-tabs-mode: nil; tab-width: 4; c-basic-offset: 4; -*-
# vim:fileencodings=utf-8:expandtab:tabstop=4:shiftwidth=4:softtabstop=4

set -x

export LANG=C LANGUAGE=C

device_node="/dev/sda"

resize_early ()
{
    total="$(fdisk -l $device_node | grep -o -E 'total (.*) sectors' | awk '{print $2}')"

    tmp="$(mktemp)"

    fdisk -l | grep ^$device_node | grep Linux$ | sed 's/*/ /' | cat > $tmp
    if read device start end blocks id system; then
        ubuntu="${device:8}"
        if fdisk -l | grep -e ^$device_node$ubuntu -e '*' 2>&1; then
            ubuntu_boot=1
        else
            ubuntu_boot=0
        fi
        ubuntu_start="$start"
        ubuntu_end="$end"
        ubuntu_size="$(expr $end - $start + 1)"
    fi < $tmp

    fdisk -l | grep ^$device_node | grep swap | cat > $tmp
    if read device start end blocks id system; then
        swap="${device:8}"
        swap_size="$(expr $end - $start + 1)"
        swap_start="$(expr $total - $swap_size)"
        swap_end="$(expr $swap_start + $swap_size - 1)"
    fi < $tmp

    rm $tmp

    if [ -n "$swap" ]; then
        parted $device_node rm 4
        swap_uuid="$(grep ^UUID.*swap /target/etc/fstab | awk '{print $1}' | cut -d '=' -f 2)"
        sed -i "s/UUID=$swap_uuid/#UUID=$swap_uuid/" /target/etc/fstab

        ubuntu_end="$(expr $swap_start - 1)"
        ubuntu_size="$(expr $ubuntu_end - $ubuntu_start + 1)"

        cat - <<-ENDLINE | fdisk $device_node || :
d
$ubuntu
d
$swap
n
p
$ubuntu
$ubuntu_start
$ubuntu_end
t
$ubuntu
83
n
p
$swap
$swap_start
$swap_end
t
$swap
82
w
ENDLINE
    else
        parted $device_node rm 3
        cat - <<-ENDLINE | fdisk $device_node || :
d
$ubuntu
n
p
$ubuntu
$ubuntu_start

t
$ubuntu
83
w
ENDLINE
    fi

    if [ "$ubuntu_boot" = 1 ]; then
        cat - <<-ENDLINE | fdisk $device_node || :
a
$ubuntu
w
ENDLINE
    fi
}

resize_late ()
{
    tmp="$(mktemp)"

    fdisk -l | grep ^$device_node | grep Linux$ | sed 's/*/ /' | cat > $tmp
    if read device start end blocks id system; then
        ubuntu="${device:8}"
    fi < $tmp

    fdisk -l | grep ^$device_node | grep swap | cat > $tmp
    if read device start end blocks id system; then
        swap="${device:8}"
    fi < $tmp

    rm $tmp

    if [ -n "$swap" ]; then
        swap_uuid="$(grep ^#UUID.*swap /etc/fstab | awk '{print $1}' | cut -d '=' -f 2)"
        if [ -n "$swap_uuid" ]; then
            sed -i "s/#UUID=$swap_uuid/UUID=$swap_uuid/" /etc/fstab
            if [ -n "$swap_uuid" ]; then
                mkswap -U $swap_uuid $device_node$swap
                swapon -a
            fi
        fi
    fi

    resize2fs $device_node$ubuntu
}

case $1 in
    (early)
        resize_early
        ;;
    (late)
        resize_late
        ;;
esac

exit 0
