#!/bin/bash

if [ $(id -u) -ne 0 ]; then
  printf "Script must be run as root. Try 'sudo raspi-config'\n"
  exit 1
fi

if [ ! -x /usr/bin/whiptail ]; then
  printf "Please install 'libnewt' package.\n"
  exit 1
fi

ASK_TO_REBOOT=0

do_info() {
  whiptail --msgbox "\
This tool provides a straight-forward way of doing initial 
configuration of the Raspberry Pi. Although it can be run 
at any time, some of the options may have difficulties if 
you have heavily customised your installation.\
" 20 70 1
}

do_expand_rootfs() {
  if [ ! -x /usr/sbin/parted ]; then
    whiptail --msgbox "Please install 'parted' package." 20 60 2
    return 0
  fi

  ROOT_PART="$(findmnt / -o source -n)"
  ROOT_DEV="/dev/$(lsblk -no pkname "$ROOT_PART")"

  PART_NUM="$(echo "$ROOT_PART" | grep -o "[[:digit:]]*$")"

  if [ "$PART_NUM" -ne 2 ]; then
    whiptail --msgbox "Your partition layout is not currently supported by this tool. You are probably using NOOBS, in which case your root filesystem is already expanded anyway." 20 60 2
    return 0
  fi

#  LAST_PART_NUM=$(parted "$ROOT_PART" -ms unit s p | tail -n 1 | cut -f 1 -d:)
#  if [ $LAST_PART_NUM -ne $PART_NUM ]; then
#    whiptail --msgbox "$ROOT_PART is not the last partition. Don't know how to expand" 20 60 2
#    return 0
#  fi

  # Get the starting offset of the root partition
  PART_START=$(parted "$ROOT_DEV" -ms unit s p | grep "^${PART_NUM}" | cut -f 2 -d: | sed 's/[^0-9]//g')
  [ "$PART_START" ] || return 1
  # Return value will likely be error for fdisk as it fails to reload the
  # partition table because the root fs is mounted
  PART_END=$(blockdev --getsz $ROOT_DEV|awk '$1=NF?$NF-1:X')
  fdisk "$ROOT_DEV" <<EOF
p
d
$PART_NUM
n
p
$PART_NUM
$PART_START
$PART_END

p
w
EOF
  ASK_TO_REBOOT=1

# now set up an init.d script

cat > /lib/systemd/system/resize2fs_once.service << EOF
[Unit]
Description=Resize2fs_once Service
[Service]
Type=idle
ExecStart=/usr/local/bin/resize2fs_once.sh
StandardOutput=null
TimeoutStartSec=infinity
[Install]
WantedBy=multi-user.target
Alias=resize2fs_once.service
EOF

# Enable the unit file so it runs after a reboot
systemctl enable resize2fs_once.service

echo "writing the resize2fs_once script"

cat  > /usr/local/bin/resize2fs_once.sh << EOF
#!/bin/sh
    resize2fs -p ${ROOT_PART}
    rm /lib/systemd/system/resize2fs_once.service
    systemctl disable resize2fs_once.service
    rm /usr/local/bin/resize2fs_once.sh
    /bin/sync
    /usr/bin/systemctl reboot
EOF

  chmod +x /usr/local/bin/resize2fs_once.sh &&
  whiptail --msgbox "Root partition has been resized.\n\
The filesystem will be enlarged upon the next reboot" 20 60 2
}

do_list_timezones() {
  timedatectl list-timezones
}

do_change_pass() {
  read -p "Enter new password for user alarm : " PASS
  echo "alarm:$PASS" | chpasswd
}

do_change_timezone() {
  read -p "Enter your timezone name : " TZONE
  timedatectl set-timezone $TZONE
  YOURTZ=$(timedatectl |grep Time | sed 's/       //g')
  echo "Updating $YOURTZ"
  sleep 3
}

do_timezone_menu() {
  FUN=$(whiptail --title "Raspberry Pi Software Configuration Tool (raspi-config)" --menu "Timezone Options" 20 80 12 --cancel-button Back --ok-button Select \
    "I1 Change List timezones" "Display a list of available timezones" \
    "I2 Change Timezone" "Set up timezone to match your location" \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    return 0
  elif [ $RET -eq 0 ]; then
    case "$FUN" in
      I1\ *) do_list_timezones ;;
      I2\ *) do_change_timezone ;;
      *) whiptail --msgbox "Programmer error: unrecognized option" 20 60 1 ;;
    esac || whiptail --msgbox "There was an error running option $FUN" 20 60 1
  fi
}

do_finish() {
  if [ $ASK_TO_REBOOT -eq 1 ]; then
    whiptail --yesno "Would you like to reboot now?" 20 60 2
    if [ $? -eq 0 ]; then # yes
      sync
      reboot
    fi
  fi
  exit 0
}

while true; do
  FUN=$(whiptail --menu "Expand SD card" 15 60 5 --cancel-button Finish --ok-button Select \
    "expand_rootfs" "Expand root partition to fill SD card" \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    do_finish
  elif [ $RET -eq 0 ]; then
    "do_$FUN" || whiptail --msgbox "There was an error running do_$FUN" 20 60 1
  else
    exit 1
  fi
done