#!/usr/bin/env bash

sectors_to_size() {
  SIZE=$(( ${1} * 512 ))
  MAG=0
  until [ $(awk "BEGIN { print ${SIZE} < 1000 }") -eq 1 ]; do
    SIZE=$(awk "BEGIN { print ${SIZE} / 1024 }" )
    ((MAG++))
  done
  case ${MAG} in
    0) TAG="" ;;
    1) TAG="Ki" ;;
    2) TAG="Mi" ;;
    3) TAG="Gi" ;;
    4) TAG="Ti" ;;
  esac

  printf "%.4g%s\n" "${SIZE}" "${TAG}B"
}

[ $EUID -gt 0 ] && echo "Must run as superuser." && exit 1

[ -t 0 ] && DEVLIST="$(fdisk -l)" || DEVLIST="$(cat -)"

for DEVICE in $(echo "${DEVLIST}" | grep '^Disk /' | grep -v '/dev/ram' | cut -d\  -f2 | cut -d: -f1); do

  [ -t 0 ] && FDISKDATA="$(fdisk -l "${DEVICE}")" || FDISKDATA="${DEVLIST}"
  COLS="$(tput cols)"
  SECTORS="$(echo "${FDISKDATA}" | grep sectors | grep -v 'sectors of' | rev | cut -d\  -f2 | rev)"
  #PARTLIST="$(echo "${FDISKDATA}" | grep '^/' | sed 's/  */ /g' | cut -d\  -f1-3)"
  PARTLIST="$(echo "${FDISKDATA}" | grep '^/' | grep -v 'Extended$' | sed -e 's/*/ /' -e 's/  \+/ /g' | cut -d\  -f1-3)"

  PARTITIONS=()
  SIZE_LOG=0
  PARTSECTS=()
  PARTSIZES=()
  SIZETOTAL=0
  PRINTCOLS=0
  PCOLS=()

  if [ -z "${PARTLIST}" ]; then
    PARTITIONS=${DEVICE}
    SIZE_LOG=1
    PARTSECTS=${SECTORS}
    PARTSIZES=1
    SIZETOTAL=1
  else
    LASTEND=0
    while read -r PARTITION PARTSTART PARTEND; do
      if [ ${PARTSTART} -gt $(( ${LASTEND} + 2049 )) ]; then
        PARTITIONS+=( "Gap" )
        SIZE_LOG=( $(awk "BEGIN { print log( ${PARTSTART} - ${LASTEND} ) / log( ${SECTORS} ) }") )
        PARTSECTS+=( $(awk "BEGIN { print ${PARTSTART} - ${LASTEND} }") )
        PARTSIZES+=( ${SIZE_LOG} )
        SIZETOTAL=$(awk "BEGIN { print ${SIZETOTAL} + ${SIZE_LOG} }")
      fi
      PARTITIONS+=( ${PARTITION} )
      SIZE_LOG=( $(awk "BEGIN { print log( ${PARTEND} - ${PARTSTART} + 1 ) / log( ${SECTORS} ) }") )
      PARTSECTS+=( $(awk "BEGIN { print ${PARTEND} - ${PARTSTART} + 1 }") )
      PARTSIZES+=( ${SIZE_LOG} )
      SIZETOTAL=$(awk "BEGIN { print ${SIZETOTAL} + ${SIZE_LOG} }")
      LASTEND=${PARTEND}
    done <<< "${PARTLIST}"
  fi

  FACTOR=$(awk "BEGIN { print ( ${COLS} - ${#PARTITIONS[@]} - 1 ) / ${SIZETOTAL} }")

  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    PRINTCOLS=$(( ${PRINTCOLS} + $(awk "BEGIN { print int ( ${PARTSIZES[${LCV}]} * ${FACTOR} + 0.5 ) }") ))
    PCOLS+=( $(awk "BEGIN { print int ( ${PARTSIZES[${LCV}]} * ${FACTOR} + 0.5 ) }") )
  done

  OVERFLOW=$(( ${PRINTCOLS} + ${LCV} + 1 - ${COLS} ))

  ####### TOP LINE #######

  printf "%s" "┌"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    [ ${LCV} = $(( ${#PARTITIONS[@]} - 1 )) ] && (( PCOLS[${LCV}]-=${OVERFLOW} ))
    printf '─%.0s' $(seq 1 ${PCOLS[${LCV}]})
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "┬"
  done
  printf "%s\n" "┐"

  ##### DEVICE NAME ######

  printf "%s" "│"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    OUTPUT=${PARTITIONS[${LCV}]##*/}
    PADL=( $(( ( ${PCOLS[${LCV}]} - ${#OUTPUT} ) / 2 )) )
    PADR=( $(( ${PCOLS[${LCV}]} - ${#OUTPUT} - ${PADL} )) )
    printf '%*s%s%*s' ${PADL} "" ${OUTPUT} ${PADR} ""
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "│"
  done
  printf "%s\n" "│"

  #### PARTITION SIZE ####

  printf "%s" "│"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    OUTPUT=$(sectors_to_size "${PARTSECTS[${LCV}]}")
    PADL=( $(( ( ${PCOLS[${LCV}]} - ${#OUTPUT} ) / 2 )) )
    PADR=( $(( ${PCOLS[${LCV}]} - ${#OUTPUT} - ${PADL} )) )
    printf '%*s%s%*s' ${PADL} "" ${OUTPUT} ${PADR} ""
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "│"
  done
  printf "%s\n" "│"

  #### PARTITION TYPE ####

  printf "%s" "│"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    OUTPUT=$(blkid -o export "${PARTITIONS[${LCV}]}" | grep '^\(TYPE=\)' | cut -d= -f2 | sed 's:linux_raid_member:raid:g;s:zfs_member:zfs:g')
    PADL=$(( ( ${PCOLS[${LCV}]} - ${#OUTPUT} ) / 2 ))
    PADR=$(( ${PCOLS[${LCV}]} - ${#OUTPUT} - ${PADL} ))
    printf '%*s%s%*s' ${PADL} "" "${OUTPUT}" ${PADR} ""
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "│"
  done
  printf "%s\n" "│"

  ### FILESYSTEM LABEL ###

  printf "%s" "│"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    OUTPUT=$(blkid -o export "${PARTITIONS[${LCV}]}" | grep '^\(LABEL=\)' | cut -d= -f2)
    PADL=$(( ( ${PCOLS[${LCV}]} - ${#OUTPUT} ) / 2 ))
    PADR=$(( ${PCOLS[${LCV}]} - ${#OUTPUT} - ${PADL} ))
    printf '%*s%s%*s' ${PADL} "" "${OUTPUT}" ${PADR} ""
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "│"
  done
  printf "%s\n" "│"

  ##### BOTTOM LINE ######

  echo -n "└"
  for (( LCV = 0; LCV < ${#PARTITIONS[@]}; LCV++ )); do
    printf '─%.0s' $(seq 1 ${PCOLS[${LCV}]})
    [ ${LCV} -lt $(( ${#PARTITIONS[@]} - 1 )) ] && echo -n "┴"
  done
  echo "┘"

done