#!/bin/sh

# Get the root mtd device number (mtdX) from "/dev/ubiblockX_Y on /"
findrootmtd() {
  rootmatch=" on / "
  m="$(mount | grep "${rootmatch}" | grep "ubiblock")"
  m="${m##*ubiblock}"
  m="${m%_*}"
  if [ -z "${m}" ]; then
    # default to bmc mtd (0)
    m=0
  fi
  echo "mtd${m}"
}

findrootubi() {
  rootmatch=" on / "
  m="$(mount | grep "${rootmatch}")"
  m="${m##*ubiblock}"
  m="${m% on*}"
  echo "ubi${m}"
}

# Get the mtd device number (mtdX)
findmtd() {
  m="$(grep -xl "$1" /sys/class/mtd/*/name)"
  m="${m%/name}"
  m="${m##*/}"
  echo "${m}"
}

# Get the ubi device number (ubiX_Y)
findubi() {
  u="$(grep -xl "$1" /sys/class/ubi/ubi?/subsystem/ubi*/name)"
  u="${u%/name}"
  u="${u##*/}"
  echo "${u}"
}

# Get all ubi devices on a specific mtd that match requested string
findubi_onmtd() {
  u="$(grep -h "$1" /sys/class/ubi/ubi"$2"/subsystem/ubi"$2"*/name)"
  u="${u%/name}"
  u="${u##*/}"
  echo "${u}"
}

# Make space on flash before creating new volumes. This can be enhanced
# determine current flash usage. For now only keep a "keepmax" number of them
ubi_remove_volumes()
{
  kernelname="$(fw_printenv -n kernelname)"
  curversion="${kernelname##*-}"
  rootubi="$(findrootubi)"

  # Just keep max number of volumes before updating, don't delete the version
  # the BMC is booted from, and when a version is identified to be deleted,
  # delete both the rofs and kernel volumes for that version.
  rmnames="$(findubi_onmtd "${name%-*}-" "${ro}")"
  rmnames=(${rmnames})
  ubicount="${#rmnames[@]}"
  while [ ${ubicount} -ge ${keepmax} ]; do
    # Loop through existing volumes and skip currently active ones
    for (( index=0; index<${#rmnames[@]}; index++ )); do
      rmname="${rmnames[${index}]}"
      rmversion="${rmname##*-}"
      rmubi="$(findubi "rofs-${rmversion}")"
      if [[ ( "${rmubi}" != "${rootubi}" ) &&
            ( "${rmname}" != "${kernelname}" ) ]]; then
        ubi_remove "rofs-${rmversion}"
        ubi_remove "kernel-${rmversion}"
        break
      fi
    done
    # Decrease count regardless to avoid an infinite loop
    (( ubicount-- ))
  done
}

ubi_rw() {
  rwmtd="$(findmtd "${reqmtd}")"
  rw="${rwmtd#mtd}"
  ubidev="/dev/ubi${rw}"

  if [ "$(fw_printenv rwreset)" == "rwreset=true" ]; then
    ubi_remove "${name}"
    fw_setenv rwreset
  fi

  # Create a ubi volume of size 4MB, that is the current size of the rwfs image
  vol="$(findubi "${name}")"
  if [ -z "${vol}" ]; then
    ubimkvol "${ubidev}" -N "${name}" -s 4MiB
  fi
}

ubi_ro() {
  keepmax=2 # Default 2 volumes per mtd
  romtd="$(findmtd "${reqmtd}")"
  romtd2="$(findmtd "${reqmtd2}")"

  if [ ! "${romtd}" == "${romtd2}" ]; then
    # Request to use alternate mtd device, choose the non-root one
    keepmax=1 # 1 volume on each of the requested mtds
    rootmtd="$(findrootmtd)"
    if [ "${rootmtd}" == "${romtd}" ]; then
      romtd="${romtd2}"
    fi
  fi
  ro="${romtd#mtd}"
  ubidev="/dev/ubi${ro}"

  ubi_remove_volumes

  # Create a static ubi volume
  # TODO Get the actual image size openbmc/openbmc#1840
  vol="$(findubi "${name}")"
  if [ -z "${vol}" ]; then
    ubimkvol "${ubidev}" -N "${name}" -s "${imgsize}" --type=static
    vol="$(findubi "${name}")"
  fi
}

# Squashfs images need a ubi block
ubi_block() {
  vol="$(findubi "${name}")"
  ubidevid="${vol#ubi}"
  block="/dev/ubiblock${ubidevid}"
  if [ ! -e "$block" ]; then
    ubiblock --create "/dev/ubi${ubidevid}"
    if [ $? != 0 ]; then
      echo "Unable to create ubiblock ${name}:${ubidevid}"
      return 1
    fi
  fi
  copy_env_var_to_alt ubiblock
  ubi_setenv "ubiblock=$(echo "${ubidevid}" | sed 's/_/,/')"
  copy_env_var_to_alt root
  ubi_setenv "root=${block}"
}

ubi_updatevol() {
  vol="$(findubi "${name}")"
  ubidevid="${vol#ubi}"
  img="/tmp/images/${version}/${imgfile}"
  ubiupdatevol "/dev/ubi${ubidevid}" "${img}"
  if [ $? != 0 ]; then
    echo "Unable to update volume ${name}!"
    return 1
  fi
  if [[ "${name}" =~ "kernel" ]]; then
    copy_env_var_to_alt kernelname
    ubi_setenv "kernelname=${name}"
  fi
}

ubi_remove() {
    rmname="$1"
    vol="$(findubi "${rmname}")"

    if [ ! -z "$vol" ]; then
        vol="${vol%_*}"

        if grep -q "${vol}:$rmname" /proc/mounts; then
            mountdir=$(grep "${vol}:$rmname" /proc/mounts | cut -d " " -f 2)
            umount "$mountdir"
            rm -r "$mountdir"
        fi

        ubirmvol "/dev/${vol}" -N "$rmname"
    fi
}

remount_ubi() {
  # Get a list of all devices formatted as UBI
  for file in /dev/mtd*; do
    if [[ $(hexdump -C -n 3 $file) =~ "UBI" ]]; then
      mtd="${file: -1}"
      if [[ $mtd =~ ^-?[0-9]+$ ]]; then
        mtds=${mtd},${mtds}
      fi
    fi
  done

  IFS=',' read -r -a mtds <<< "$mtds"
  mtds=($(echo "${mtds[@]}" | tr ' ' '\n' | sort -u | tr '\n' ' '))
  for mtd in ${mtds[@]}; do
    # Re-attach mtd device to ubi if not already done
    ubiattach /dev/ubi_ctrl -m "${mtd}" -d "${mtd}" &> /dev/null
    # Get information on all ubi volumes
    ubinfo=$(ubinfo -d ${mtd})
    presentVolumes=${ubinfo##*:}
    IFS=', ' read -r -a array <<< "$presentVolumes"
    for element in ${array[@]}; do
      elementProperties=$(ubinfo -d $mtd -n $element)
      # Get ubi volume name by getting rid of additional properties
      name=${elementProperties#*Name:}
      name="${name%Character*}"
      name="$(echo -e "${name}" | tr -d '[:space:]')"

      if [[ ${name} == rofs-* ]]; then
        mountdir="/media/${name}"
        mkdir -p "${mountdir}"
        ubiblock --create /dev/ubi${mtd}_${element} &> /dev/null
        mount -t squashfs -o ro "/dev/ubiblock${mtd}_${element}" "${mountdir}"
      fi
    done
  done
}

# Read the current env variable and set it on the alternate boot env
copy_env_var_to_alt() {
  varName=$1
  value="$(fw_printenv -n "${varName}")"
  fw_setenv -c /etc/alt_fw_env.config "${varName}" "${value}"
}

ubi_setenv() {
    variable=$1
    varName="${variable%=*}"
    value="${variable##*=}"
    fw_setenv "$varName" "$value"
}

mtd_write() {
  flashmtd="$(findmtd "${reqmtd}")"
  img="/tmp/images/${version}/${imgfile}"
  flashcp -v ${img} /dev/${flashmtd}
}

case "$1" in
  mtduboot)
    reqmtd="$2"
    version="$3"
    imgfile="image-u-boot"
    mtd_write
    ;;
  ubirw)
    reqmtd="$2"
    name="$3"
    ubi_rw
    ;;
  ubiro)
    reqmtd="$(echo "$2" | cut -d "+" -f 1)"
    reqmtd2="$(echo "$2" | cut -d "+" -f 2)"
    name="$3"
    version="$4"
    imgfile="image-rofs"
    imgsize="16MiB"
    ubi_ro
    ubi_updatevol
    ubi_block
    ;;
  ubikernel)
    reqmtd="$(echo "$2" | cut -d "+" -f 1)"
    reqmtd2="$(echo "$2" | cut -d "+" -f 2)"
    name="$3"
    version="$4"
    imgfile="image-kernel"
    imgsize="4MiB"
    ubi_ro
    ubi_updatevol
    ;;
  ubiremove)
    name="$2"
    ubi_remove "${name}"
    ;;
  ubisetenv)
    ubi_setenv "$2"
    ;;
  ubiremount)
    remount_ubi
    ;;
  *)
    echo "Invalid argument"
    exit 1
    ;;
esac
rc=$?
if [ ${rc} -ne 0 ]; then
  echo "$0: error ${rc}"
  exit ${rc}
fi