#!/bin/bash

while [ "${1:-}" != '' ]; do
  case "$1" in
  '--default-bootcfg')               default_bootcfg=1    ;;
  '--fip2boot')                      fip2boot=1           ;;
  '--boot2fip')                      boot2fip=1           ;;
  '--download2root')                 download2root=1      ;;
  '--nand-force-erase')              nand_force_erase=1   ;;
  '--nand-format')                   nand_format=1        ;;
  '--nand-image')                    nand_image=1         ;;
  '--nand-update')                   nand_update=1        ;;
  '--write2dtb' | '--dtb')           write2dtb=1          ;;
  '--write2atf' | '--atf')           write2atf=1          ;;
  '--write2fip' | '--fip')           write2fip=1          ;;
  '--write2extlinux' | '--extlinux') write2extlinux=1     ;;
  '--uboot-install')                 uboot_install=1      ;;
  '--uboot-remove')                  uboot_remove=1       ;;
  '--uartboot')                      uartboot=1           ;;
  '--remove-dtb')                    remove_dtb=1         ;;
  '--clean-boot')                    clean_boot=1         ;;
  '--create-parts')                  create_parts=1       ;;
  '--target')                 shift; target="${1}"        ;;
  '--target'=*)                      target="${1#*=}"     ;;
  '--device')                 shift; device="${1}"        ;;
  '--device'=*)                      device="${1#*=}"     ;;
  '--pkgbase')                shift; pkgbase="${1}"       ;;
  '--pkgbase'=*)                     pkgbase="${1#*=}"    ;;
  '--set-atf-linuxpkg')              set_atf_linuxpkg=1   ;;
  *) [ -n "$1" ] && echo "Unknown option: $1"
     cat <<-EOF
	Usage: $(basename "$0") [OPTION]...
	  --default-bootcfg        Restore default bootcfg, adds --write2fip
	  --fip2boot               Convert fip partition to boot partition bootchain (sd/emmc)
	  --boot2fip               Convert boot partition to fip partition bootchain (sd/emmc)
	  --download2root          Download files needed for nand-image (when started from initrd)
	  --nand-force-erase       Force erase nand, including bad blocks and wear history
	  --nand-format            Format the nand, also runs update
	  --nand-image             Create nand image, also runs update
	  --nand-update            Updates all files on nand, only writes when needed
	  --write2dtb              Combine dtbos with dtb and create one dtb file
	  --write2atf              Write arm-trusted-firmware
	  --write2fip              Create all files needed for fip and write it, adds --write2dtb
	  --write2extlinux         Create a new /boot/extlinux/extlinux.conf
	  --uboot-install          Copies U-Boot to /boot/u-boot.bin (writes to fip if necessary),
	                             also creates /boot/extlinux/extlinux.conf if not present
	  --uboot-remove           Removes /boot/u-boot.bin
	  --uartboot               Create a uartboot image
	  --pkgbase ...            Specify linuxpkg to create files for
	  --set-atf-linuxpkg       Set linuxpkg atf will directly boot, specified in pkgbase
	  --remove-dtb             Remove dtb file
	EOF
    exit 1
    ;;
  esac
  shift
done

extlinux="/boot/extlinux/extlinux.conf"

[ -n "$set_atf_linuxpkg" ] && [ -n "$pkgbase" ] && default_bootcfg=1
[ -n "$uboot_install" ]    && write2fip=1
[ -n "$uboot_remove" ]     && write2fip=1
[ -n "$default_bootcfg" ]  && write2fip=1
[ -n "$set_atf_linuxpkg" ] && write2fip=1
[ -n "$write2fip" ]        && write2dtb=1
[ -n "$nand_image" ]       && nand_update=1
[ -n "$nand_format" ]      && nand_update=1
[ -n "$uboot_install" ]    && [ ! -f "$extlinux" ] && write2extlinux=1
[ -n "$create_parts" ]     && write2atf=1
[ -n "$create_parts" ]     && [ ! -n "$boot2fip" ] && fip2boot=1 # bootchain with boot partition unless specified otherwise
[ -n "$fip2boot" ]         && write2fip=1
[ -n "$boot2fip" ]         && write2fip=1

#############
# Functions #
#############
function checkbootconf {
  local conf="$(cat "/boot/bootcfg/${1}" 2>/dev/null)"
  [ -z "${conf}" ] && return
  [ -f "${conf}" ] && return
  echo "Error: ${conf} does not exist, setting /boot/bootcfg/${1} to default value."
  rm -vf "/boot/bootcfg/${1}"
}

function getconf {
  local path
  [[ "$1" == "boot" ]] && path="/boot/bootcfg" || path="/etc/rootcfg"
  local conf=$(cat "${path}/${3}" 2>/dev/null)
  if [ -z "${conf}" ]; then
    conf="${4}"
    [[ "${2}" == "w" ]] && echo -n "${conf}" >"${path}/${3}"
  fi
  echo "${conf}"
}

function checkextlinux {
  [ ! -f "${1}" ] && return
  local option file
  for option in LINUX INITRD FDT FDTOVERLAYS; do
    files=$(grep -iwF ${option} "${1}"| xargs -L1 echo)
    [[ "${option}" != "$(echo ${files} | cut -d' ' -f1)" ]] && continue
    for files in $(echo ${files} | cut -d' ' -f2-); do
      for file in ${files}; do
        file="$(dirname ${1})/..${file/'/boot/'/'/'}"
        [ ! -f "${file}" ] && echo -e "\nWARNING: extlinux.conf ${option} $(realpath ${file}) missing! :WARNING\n"
      done
    done
  done
}

function parts {
  lsblk "$1" -lnpo name 2>/dev/null
}

function getdev {
  local dev="$(blkid -t "PARTLABEL=${target}-${device}-$1" -l -o device)"
  if [ -n "${dev}" ]; then echo "${dev}"; return; fi
  [ -z "${pkatfdevice}" ] && return
  dev="$(blkid -t "PARTLABEL=$1" -o device $(parts "/dev/${pkatfdevice}"))"
  [ -n "${dev}" ] && echo "${dev}"
}

function fix_partlabelfipboot {
  if grep "${target}-${device}-" "${atffile}" 2>/dev/null; then
    echo "This atf uses extended partition name"
    if [[ "${partlabelfipboot}" != "${target}-${device}-${typefipboot}" ]]; then
      echo "Changing partlabel $fipbootdevice to ${target}-${device}-${typefipboot}"
      parted -s -- "/dev/${pkfipboot}" name ${nrfipboot} ${target}-${device}-${typefipboot}
      sync
    fi
  else
    echo "This atf uses standard 'fip' or 'boot' partition name"
    if [[ "${partlabelfipboot}" != "${typefipboot}" ]]; then
      echo "Changing partlabel ${fipbootdevice} to ${typefipboot}"
      parted -s -- "/dev/${pkfipboot}" name ${nrfipboot} ${typefipboot}
      sync
    fi
  fi
}

function imageid_noarchive {
  od --skip-bytes=$((4*14)) --read-bytes=4 --address-radix=n -x "${1}" \
    | sed 's/ //g'
}

function imageid {
  if [[ "$(file -b --mime-type ${1})" == "application/x-xz" ]]; then
    zstd -dcf --format=xz "${1}" | imageid_noarchive -
  else
    imageid_noarchive "${1}"
  fi
}

function get_emmcbootpart {
  local bootpart="0"$(mmc extcsd read ${1} | grep PARTITION_CONFIG | cut -f 2 -d'0' | cut -f 1 -d']')
  echo $(( (${bootpart} & 0x38) >> 3 ))
}

function set_emmcbootpart {
  [ ${2} -eq $(get_emmcbootpart ${1}) ] && return
  mmc bootpart enable ${2} 1 ${1}
}

function fix_emmcbootpart {
  if [ -b "${1}boot0" ]; then # Writing to EMMC boot device
    local hdr=$(head -c 4 "${atfdevice}")
    if [[ "${hdr}" == "SDMM" ]] || [[ "${hdr}" == "BRLY" ]]; then
      echo "Customised ATF able to boot from ${1}"
      set_emmcbootpart ${1} 7
    else
      echo -e "ATF needs boot from ${1}boot0, copying..."
      echo 0 > /sys/block/${1/'/dev/'/}boot0/force_ro
      dd bs=4M if="${atfdevice}" of="${1}boot0" conv=fsync
      echo 1 > /sys/block/${1/'/dev/'/}boot0/force_ro
      set_emmcbootpart ${1} 1
    fi
  fi
}

function mtdnr {
  cat /proc/mtd | grep '"'$1'"' | cut -d':' -f1 | tr -d [:alpha:]
}

function ubivol {
  for u in /sys/class/ubi/${ubidev/"/dev/"/}_*/name; do
    if [[ "$(cat ${u})" == "$1" ]]; then
      echo $(basename $(dirname "${u}"))
      return
    fi
  done
}

function ubiupdate {
  if ! diff -q "$1" "$2"; then
    echo "Updating "$(basename "$1")" on NAND:"
    cp -vf "$1" "$2"
  else
    echo "Skipping "$(basename "$1")" on NAND"
  fi
}

function get_ubidev {
  local ubidevice
  ubidevice="$(basename /sys/class/mtd/mtd${1}/ubi*)"
  [ ! -e "/sys/class/ubi/${ubidevice}" ] && return 1
  echo "/dev/${ubidevice}"
  return 0
}

function _ubimkvol {
  if [[ "${1}" == "image" ]]; then
    local OPTIND=2  n s t N m option
    while getopts ":n:s:t:N:m" option; do
      case "${option}" in
        n) n="${OPTARG}";;
        s) s="${OPTARG}";;
        m) m="true";;
        t) t="${OPTARG}";;
        N) N="${OPTARG}";;
      esac
    done
    (
      echo -e "[${N}-volume]\nmode=ubi\nimage=${tmp}/${N}.bin\nvol_id=${i}\nvol_name=${N}\nvol_alignment=1"
      [[ ${m} == "true" ]]   && echo "vol_flags=autoresize" || echo "vol_size=${s}"
      [[ ${t} == "static" ]] && echo "vol_type=static"      || echo "vol_type=dynamic"
      echo
    ) >> ${tmp}/ubinize.conf
  else
    ubimkvol $@
  fi
}

function write_extlinux {
  local pkgb
  mkdir -p $(dirname "$1")
  cat <<-EOF | tee "$1"
	DEFAULT ${linuxpkg}
	  MENU TITLE U-Boot menu
	  PROMPT 0
	  TIMEOUT 30
	EOF
  for pkgb in $(shopt -s nullglob; for f in /lib/modules/*/pkgbase ; do cat $f; echo; done | xargs); do
    cat <<-EOF | tee "$1"
	LABEL ${pkgb}
	  MENU LABEL Archlinux ARM for BananaPi Routers
	  LINUX ${2/${linuxpkg}/${pkgb}}
	  INITRD ${3/${linuxpkg}/${pkgb}}
	  FDT ${4/${linuxpkg}/${pkgb}}
	EOF
  done
}

function write_dtb {
  cp -vf "$3" "${tmp}/fixed.dtb"
  rm -rf "${tmp}/dtbos/"* 
  for dts in "${@:4}"; do
    [[ "${dts}" =~ '*' ]] && continue
    dtsname=$(basename ${dts})
    echo "Creating .dtbo from ${dtsname}"
    cat "${dts}" | grep -v -e '^#define' > "${tmp}/dtbos/${dtsname}"
    cat "${dts}" | grep -e '^#define' | while read -r line ; do
      macroname=$(echo "${line}" | tr -s ' \t' ' ' | cut -d ' ' -f 2)
      macrotext=$(echo "${line}" | tr -s ' \t' ' ' | cut -d ' ' -f 3)
      if [ -n "${macroname}" ]; then
        echo "Applying #define ${macroname} = ${macrotext}"
        sed -i "s/${macroname}/${macrotext}/g" "${tmp}/dtbos/${dtsname}"
      fi
    done
    cat "${dts}" | grep "//fdtput" | while read -r line ; do
      echo fdtput "${tmp}/fixed.dtb" ${line/"//fdtput"/""}
           fdtput "${tmp}/fixed.dtb" ${line/"//fdtput"/""}
    done
    dtc -@ -q -I dts -O dtb -o "${tmp}/dtbos/${dtsname/.dts/.dtbo}" \
        "${tmp}/dtbos/${dtsname}"
  done
  fdtoverlay -vi "${tmp}/fixed.dtb" -o "$1" \
                  ${tmp}/dtbos/*.dtbo
  origargs="$(fdtget -ts "$1" "/chosen" "bootargs")"
  bootargs="$2 ${origargs} ${cmdline}"
  echo BOOTARGS = "${bootargs}"
  fdtput -ts "$1" "/chosen" "bootargs" "${bootargs}"
  fdtput -ts "$1" "/memory" "device_type" "memory"
  d1=$(printf "%x" $((ddrsize >> 2)) )
  d2=$(printf "%x" $(( $((ddrsize << 2)) & 15 )) )
  fdtput -tx "$1" "/memory" "reg" 00 40000000 0${d1} ${d2}0000000
  if [ -f "${initrd}" ]; then
    ins="0x48000000"
    ine="0x$(printf '%x\n' $(( ${ins} + $(du -b ${initrd} | cut -f1) )))"
    fdtput -tx "$1" "/chosen" "linux,initrd-end" "${ine}"
    fdtput -tx "$1" "/chosen" "linux,initrd-start" "${ins}"
  fi
}

function download {
  echo "Download $1..."
  pkg=$(cat "${tmp}/ericwoud/$1-"*"/desc" 2>/dev/null | grep "%FILENAME%" -A1 | head -n 2 | tail -n1)
  [ -z "${pkg}" ] && return 1
  until curl -L "${repo}"'/'"${pkg}" | xz -dc - | tar x -mC /
  do sleep 2; done
  return 0
}

function finddtbfile {
  local                 file="/boot/dtbs/${1}/mediatek/${2}"
  [ ! -f "${file}" ] && file="/boot/dtbs/${1}/${2}"
  [ ! -f "${file}" ] && file="/boot/dtbs/mediatek/${2}"
  [ ! -f "${file}" ] && file="/boot/dtbs/${2}"
  [ ! -f "${file}" ] && return
  echo "${file}"
}

#################
# Set variables #
#################
tmp="/tmp/bpir-toolbox-tmp"
rm -rf "${tmp}"
mkdir -p "${tmp}/dtbos"
mkdir -p /boot/bootcfg
mkdir -p /etc/rootcfg
mv -vf /boot/bootcfg/ddrsize /etc/rootcfg/ddrsize 2>/dev/null

rootdev="$(lsblk -pilno name,type,mountpoint 2>/dev/null | grep -G 'part /$' | head -n1 | cut -d " " -f1)"

[ -z "${target}" ] && target="$(getconf root r target)"
[ -z "${device}" ] && device="$(getconf root r device)"

if [ -z "${target}" ] || [ -z "${device}" ]; then
  partlabelroot="$(blkid "${rootdev}" -s PARTLABEL -o value)"
  [ -z "${partlabelroot}" ] && exit 1
  target="$(echo ${partlabelroot} | cut -d'-' -f1)"
  device="$(echo ${partlabelroot} | cut -d'-' -f2)"
fi

# Write target and device if not written yet
echo "Target=$(getconf root w target ${target})"
echo "Device=$(getconf root w device ${device})"

case ${target} in
  bpir64)
    default_linuxpkg="linux-bpir-git"
    default_dtb="mt7622-bananapi-bpi-r64"
    default_ddrsize=1
    nandtype="snand"
    ;;
  bpir3)
    default_linuxpkg="linux-bpir-git"
    default_dtb="mt7986a-bananapi-bpi-r3"
    default_ddrsize=2
    nandtype="spim-nand"
    ;;
  bpir3m)
    default_linuxpkg="linux-bpir-git"
    default_dtb="mt7986a-bananapi-bpi-r3-mini"
    default_ddrsize=2
    nandtype="spim-nand"
    ;;
  bpir4)
    default_linuxpkg="linux-bpir4-git"
    default_dtb="mt7988a-bananapi-bpi-r4"
    default_ddrsize=4
    nandtype="spim-nand"
    ;;
  *)
    echo "Unknown target ${target}"
    exit 1
    ;;
esac

case ${device} in
  nvme|sata)  [ -n "$create_parts" ] && write2extlinux=1 ;;
esac

default_cmdline="console=ttyS0,115200 debug=7 rw rootwait audit=0"

if [ -z "${rootdev}" ]; then
  if [ -f "/etc/bpir-is-initrd" ]; then  # Writing NAND while running from initrd
    device="${nandtype}"
    if  [ -n "$nand_update" ] || [ -n "$download2root" ]; then
      mkdir -p "${tmp}/ericwoud"
      repo="ftp://ftp.woudstra.mywire.org/repo/aarch64"
      until curl -L "${repo}/ericwoud.db" | tar -xz -mC "${tmp}/ericwoud"
      do sleep 2; done
      download "${default_linuxpkg}"
      download mkinitcpio-bpir
      download bpir-atf-git
      download bpir-uboot-git
      echo -n "${default_linuxpkg}" >"$(echo /lib/modules/*/kernel | head -n1 | xargs dirname)/pkgbase"
      bpir-initrd -P
    fi
  fi
fi

if [ -n "$default_bootcfg" ]; then
  find "/boot/bootcfg" -type f ! -name 'ddrsize' -delete
fi

if [ -n "$set_atf_linuxpkg" ] && [ -n "${pkgbase}" ]; then
  # bootcfg already cleared
  default_linuxpkg="${pkgbase}"
fi

if [ -z "$(grep /lib/modules/*/pkgbase -l -e ${default_linuxpkg})" ]; then
  # default_linuxpkg not valid, set first (alfabetically sorted) linux kernel installed
  default_linuxpkg="$(cat $(echo /lib/modules/*/kernel | head -n1 | xargs dirname)/pkgbase)"
  [ -z "${default_linuxpkg}" ] && exit 1
fi

[ -n "${pkgbase}" ] && default_linuxpkg="${pkgbase}"
linuxpkg=$(getconf root w linuxpkg "${default_linuxpkg}")
[ -z "${pkgbase}" ] && pkgbase="${linuxpkg}"

unamer="$(grep /lib/modules/*/pkgbase -l -e "${pkgbase}" | head -n1 | xargs dirname | xargs basename)"

if [ ! -f "/boot/Image" ] || [ -f "/boot/Image-${pkgbase}" ]; then
  default_linux="/boot/Image-${pkgbase}"
  default_initrd="/boot/initramfs-${pkgbase}.img"
else
  default_linux="/boot/Image"
  default_initrd="/boot/initramfs-bpir.img"
fi

checkextlinux "${extlinux}"

checkbootconf linux
checkbootconf initrd
checkbootconf atfdtb

dtb=$(getconf root r dtb "${default_dtb}")
dtbfile=$(finddtbfile "${pkgbase}" "${dtb}.dtb")
[ ! -f "${dtbfile}" ] && exit 1
default_atfdtb="${dtbfile%.dtb}-atf.dtb"

linux="$(   getconf boot w linux    "${default_linux}")"
cmdline="$( getconf boot w cmdline  "${default_cmdline}")"
initrd="$(  getconf boot w initrd   "${default_initrd}")"
atfdtb="$(  getconf boot w atfdtb   "${default_atfdtb}")"
ddrsize="$( getconf root r ddrsize  "${default_ddrsize}")"

extra=""
[[ "${ddrsize}" != "${default_ddrsize}" ]] && [[ ! "${ddrsize}" =~ "default" ]] && extra+="-${ddrsize}gb"
[ -d "/usr/share/bpir-atf" ] && atfdir="/usr/share/bpir-atf" || atfdir="/boot"
headerfile="${atfdir}/${target}-atf-${device}-header${extra}.bin"
atffile="${atfdir}/${target}-atf-${device}-atf${extra}.bin"
bl31file="${atfdir}/${target}-atf-${device}-bl31${extra}.bin"
 
if [[ "$(cat /etc/rootcfg/atffile 2>/dev/null | xargs)" != "${atffile}" ]]; then
  echo -n "${atffile}" >/etc/rootcfg/atffile
fi

# Copy DTBO's if there are none.
if [ ! -d "/boot/dtbos/" ]; then
  mkdir -p "/boot/dtbos/"
  cp -vf "/usr/share/buildR64arch/boot/${target^^}/"*             "/boot/dtbos/" 2>/dev/null
  cp -vf "/usr/share/buildR64arch/boot/${target^^}-${device^^}/"* "/boot/dtbos/" 2>/dev/null
fi

####################
# Force erase NAND #
####################
if [ -n "$nand_force_erase" ]; then
  [[ "$(tr -d '\0' 2>/dev/null </proc/device-tree/compatible)" != "bananapi,"* ]] && exit 1
  ubinr=$(mtdnr ubi)
  ubidetach -p /dev/mtd${ubinr} 2>/dev/null
  echo Y >/sys/kernel/debug/mtd/expert_analysis_mode
  for i in $(cat /proc/mtd | tail -n+2 | cut -d' ' -f1 | tr -d ':'); do
    flash_erase --noskipbad "/dev/${i}" 0 0
  done
  echo N >/sys/kernel/debug/mtd/expert_analysis_mode
fi

###############
# Format NAND #
###############
ubidev=""
if [ -n "$nand_image" ]; then
  ubidev="image"
elif [ -n "$nand_format" ]; then
  [[ "$(tr -d '\0' 2>/dev/null </proc/device-tree/compatible)" != "bananapi,"* ]] && exit 1
  ubinr=$(mtdnr ubi)
  ubidetach -p /dev/mtd${ubinr} 2>/dev/null
  ubiformat -y /dev/mtd${ubinr}
  ubiattach -p /dev/mtd${ubinr}
  ubidev=$(get_ubidev ${ubinr})
  [[ $? != 0 ]] && exit 1
  [[ "${ubidev}" == "/dev/ubi_ctrl" ]] && exit 1
fi
if [ -n "${ubidev}" ]; then
  rm -f "${tmp}/ubinize.conf"
  i=0
  _ubimkvol "${ubidev}" -n ${i} -N fip       -s 2MiB   -t static
  ((i++))
  _ubimkvol "${ubidev}" -n ${i} -N ubootenv  -s 128KiB
  ((i++))
  _ubimkvol "${ubidev}" -n ${i} -N ubootenv2 -s 128KiB
  ((i++))
  _ubimkvol "${ubidev}" -n ${i} -N rootfs    -m
  if [ "${ubidev}" != "image" ]; then
    while [ ! -e "${ubidev}_${i}" ]; do sleep 0.1; done
    mkfs.ubifs "${ubidev}_${i}"
  fi
fi

###############
# Update NAND #
###############
if [ -n "$nand_update" ]; then
  nandbl31file="/usr/share/bpir-atf/${target}-atf-${nandtype}-bl31${extra}.bin"
  atfbin="/usr/share/bpir-atf/${target}-atf-${nandtype}-atf${extra}.bin"
  [ -f "${linux}.gz" ] && ubootlinux="${linux}.gz" || ubootlinux="${linux}"
  ubootfile="/usr/share/bpir-uboot/u-boot-${target}.bin"
  [ ! -f "${ubootfile}" ] && exit 1
  fipfiles=""
  [ -f "${nandbl31file}" ] && fipfiles+=" --soc-fw ${nandbl31file}"
  fipfiles+=" --nt-fw ${ubootfile}"
  fiptool --verbose create "${tmp}/fip.bin" ${fipfiles}
  fiptool info "${tmp}/fip.bin"
  write_dtb "${tmp}/atf.dtb" \
            "root=" \
            "${dtbfile}" \
	    "/usr/share/buildR64arch/boot/${target^^}/"*".dts" \
	    "/usr/share/buildR64arch/boot/${target^^}-NAND/"*".dts"
  write_extlinux "${tmp}/extlinux.conf" \
                 "${ubootlinux}" \
                 "${initrd}" \
                 "${atfdtb}"
  mkdir -p "${tmp}/mnt"
  if [ -n "$nand_image" ]; then
    mkdir -p "${tmp}/mnt/boot/extlinux" "${tmp}/mnt/boot/dtbs" "$(dirname "${tmp}/mnt${atfdtb}")"
    cp -vf "${tmp}/atf.dtb" "${tmp}/mnt${atfdtb}"
    cp -vf "${tmp}/extlinux.conf" "${tmp}/mnt/boot/extlinux/extlinux.conf"
    cp -vf "${ubootlinux}" "${tmp}/mnt${ubootlinux}"
    cp -vf "${initrd}" "${tmp}/mnt${initrd}"
    mkfs.ubifs -m 2048 -e 124KiB -c 800 -r "${tmp}/mnt" "${tmp}/rootfs.bin"
    echo > "${tmp}/ubootenv.bin"
    echo > "${tmp}/ubootenv2.bin"
    mkdir -p "/tmp/nandimage"
    cp -vf "${atfdir}/${target}-atf-${nandtype}-atf${extra}.bin" \
             "/tmp/nandimage/nand-${target}-bl2${extra}.bin" 2>/dev/null
    nandfile="/tmp/nandimage/nand-${target}-ubi${extra}.bin"
    ubinize -o "${nandfile}" -m 2048 -p 128KiB "${tmp}/ubinize.conf"
    ls -l "/tmp/nandimage"
  else
    [[ "$(tr -d '\0' 2>/dev/null </proc/device-tree/compatible)" != "bananapi,"* ]] && exit 1
    bl2nr=$(mtdnr bl2)
    ubinr=$(mtdnr ubi)
    ubidev=$(get_ubidev ${ubinr})
    if [[ $? != 0 ]]; then
      ubiattach -p /dev/mtd${ubinr}
      ubidev=$(get_ubidev ${ubinr})
      [[ $? != 0 ]] && exit 1
    fi
    if [[ "${ubidev}" == "/dev/ubi_ctrl" ]] || [[ "${ubidev}" == "/dev/ubi*" ]]; then exit 1; fi
    ubifip=$(ubivol fip)
    ubirootfs=$(ubivol rootfs)
    dd if=/dev/mtdblock${bl2nr} of="${tmp}/dump.bin" bs=$(du -b "${atfbin}" | cut -f1) count=1 >/dev/null 2>&1
    #nanddump -nf "${tmp}/dump.bin" -l $(du -b "${atfbin}" | cut -f1) /dev/mtd${bl2nr}
    if ! diff "${tmp}/dump.bin" "${atfbin}"; then
      echo "Updating bl2 on NAND:"
      dd if="${atfbin}" of="/dev/mtdblock${bl2nr}"
      #flashcp -v "${atfbin}" /dev/mtd${bl2nr}
    else
      echo "Skipping bl2 on NAND"
    fi
    dd if="/dev/${ubifip}" of="${tmp}/dump.bin"
    if ! diff "${tmp}/dump.bin" "${tmp}/fip.bin"; then
      echo "Updating fip on NAND"
      ubiupdatevol "/dev/${ubifip}" "${tmp}/fip.bin"
    else
      echo "Skipping fip on NAND"
    fi
    mount -t ubifs "${ubirootfs}" "${tmp}/mnt"
    while ! mountpoint -q "${tmp}/mnt"; do sleep 0.1; done
    mkdir -p ${tmp}/mnt/boot/extlinux "${tmp}/mnt/boot/dtbs"
    ubiupdate "${tmp}/atf.dtb" "${tmp}/mnt${atfdtb}"
    ubiupdate "${tmp}/extlinux.conf" "${tmp}/mnt/boot/extlinux/extlinux.conf"
    ubiupdate "${ubootlinux}" "${tmp}/mnt${ubootlinux}"
    ubiupdate "${initrd}" "${tmp}/mnt${initrd}"
    sync
    ls -lR "${tmp}/mnt"
    sync
    while mountpoint -q "${tmp}/mnt"; do umount "${tmp}/mnt"; sleep 0.1; done
  fi
fi

############
# UARTBOOT #
############
if [ -n "$uartboot" ]; then
  mkdir -p "/tmp/uartboot"
  rm -rf "/tmp/uartboot/"*
  write_dtb "${tmp}/atf.dtb" \
            "root=" \
            "${dtbfile}" \
	    "/usr/share/buildR64arch/boot/${target^^}/"*".dts" \
	    "/usr/share/buildR64arch/boot/${target^^}-EMMC/"*".dts"
  cp -vf "${atfdir}/${target}-atf-ram-atf${extra}.bin" \
          "/tmp/uartboot/uart-${target}-atf${extra}.bin" 2>/dev/null
  fipfile="/tmp/uartboot/uart-${target}-fip${extra}.bin"
  fipfiles=""
  uartbl31file="${atfdir}/${target}-atf-ram-bl31${extra}.bin"
  [ -f "${uartbl31file}" ] && fipfiles+=" --soc-fw ${uartbl31file}"
  fipfiles+=" --nt-fw ${linux}"
  fipfiles+=" --nt-fw-config ${tmp}/atf.dtb"
  fipfiles+=" --tos-fw-extra2 ${initrd}"
  fiptool --verbose create "${fipfile}" ${fipfiles}
  fiptool info "${fipfile}"
  ls -l "/tmp/uartboot"
fi

############
# EXTLINUX #
############
if [ -n "$write2extlinux" ]; then
  if mountpoint -q /boot/; then skip=5; else skip=0; fi
  write_extlinux "${extlinux}" \
                 "${linux:$skip}" \
                 "${initrd:$skip}" \
                 "${atfdtb:$skip}"
fi

##############
# U-BOOT PKG #
##############
if [ -n "$uboot_install" ]; then
  cp -vf /usr/share/bpir-uboot/u-boot-${target}-${device}.bin /boot/u-boot.bin # This is the file that boots
fi

if [ -n "$uboot_remove" ]; then
  rm -vf "/boot/u-boot.bin"
  rm -vf "${extlinux}"
  rmdir -v $(dirname "${extlinux}")
fi

#####################
# Create single dtb #
#####################
if [ -n "$write2dtb" ]; then
  rm -v "/boot/dtbos/"*".dtbo" "/boot/dtbs/"*"-fixed.dtb" 2>/dev/null # not used anymore
  write_dtb "${atfdtb/${linuxpkg}/${pkgbase}}" \
            "root=PARTLABEL=${target}-${device}-root" \
            "${dtbfile}" \
            "/boot/dtbos/"*".dts"
fi

#####################
# Remove single dtb #
#####################
if [ -n "$remove_dtb" ]; then
  rm -vf "${atfdtb/${linuxpkg}/${pkgbase}}"
fi

#####################
# Remove boot files #
#####################
if [ -n "$clean_boot" ]; then
  rm -vf  "/boot/Image-${pkgbase}"         2>/dev/null
  rm -vf  "/boot/Image-${pkgbase}.gz"      2>/dev/null
  rm -vf  "/boot/config-${unamer}"         2>/dev/null
  rm -vf  "/boot/System.map-${unamer}"     2>/dev/null
  rm -vf  "/boot/"*"${unamer}.itb"         2>/dev/null
  rm -vfr "/boot/dtbs/${pkgbase}"          2>/dev/null
  rm -vfr "/boot/dtbs/mediatek/${pkgbase}" 2>/dev/null
fi



###########################################
# Functions below need fip/boot partition #
###########################################
if [ -z "${rootdev}" ] && [ ! -f "/etc/bpir-is-initrd" ]; then
  if [ -n "$create_parts" ]; then
    echo "Root filesystem not mounted on block device."
  fi
  exit 0
fi
pkrootdev=$(lsblk -rno pkname "${rootdev}")
[ -z "${pkrootdev}" ] && exit 1
nrrootdev=$(cat "/sys/block/${pkrootdev}/$(basename ${rootdev})/partition" 2>/dev/null)
[ -z "${nrrootdev}" ] && exit 1
strootdev=$(cat "/sys/block/${pkrootdev}/$(basename ${rootdev})/start" 2>/dev/null)
[ -z "${strootdev}" ] && exit 1

if [ -n "$create_parts" ]; then
  parted -s -- "/dev/${pkrootdev}" name ${nrrootdev} "${target}-${device}-root"
  sync
fi

if [[ "${device}" != "sdmmc" ]] && [[ "${device}" != "emmc" ]]; then
  # When root is mounted on blkdev, but target is nvme or sata,
  #  then only set boot flag on rootdev and exit
  #  for sdmmc and emmc boot flag on rootdev is set with --fip2boot
  if [ -n "$create_parts" ]; then
    if mountpoint -q /boot/; then
      parted -s -- "/dev/${pkrootdev}" set ${nrrootdev} boot off
    else
      parted -s -- "/dev/${pkrootdev}" set ${nrrootdev} boot on
    fi
    sync
    parted -s -- "/dev/${pkrootdev}" print
  fi
  exit 1
fi

ATF_END_MB=1
atfdevice=$(getdev atf)
if [ -z "${atfdevice}" ]; then
  if [ -n "$create_parts" ]; then
    parted -s -- "/dev/${pkrootdev}" unit MiB mkpart "${target}-${device}-atf" 34s "${ATF_END_MB}"
    sync
    atfdevice="$(getdev atf)"
  fi
  [ -z "${atfdevice}" ] && exit 1
fi
pkatfdevice=$(lsblk -rno pkname ${atfdevice} 2>/dev/null)
[ -z "${pkatfdevice}" ] && exit 1
nratfdevice=$(cat "/sys/block/${pkatfdevice}/$(basename ${atfdevice})/partition" 2>/dev/null)
[ -z "${nratfdevice}" ] && exit 1
if [ -n "$create_parts" ]; then
    parted -s -- "/dev/${pkatfdevice}" set ${nratfdevice} legacy_boot on
    sync
fi

fipbootdevice="$(getdev fip)"
[ -z "${fipbootdevice}" ] && fipbootdevice="$(getdev boot)"
if [ -z "${fipbootdevice}" ]; then
  if [ -n "$create_parts" ]; then
    parted -s -- "/dev/${pkrootdev}" unit MiB mkpart fip "${ATF_END_MB}" "$((strootdev-1))s"
    sync
    fipbootdevice="$(getdev fip)"
  fi
  [ -z "${fipbootdevice}" ] && exit 1
fi

pkfipboot=$(lsblk -rno pkname ${fipbootdevice} 2>/dev/null)
[ -z "${pkfipboot}" ] && exit 1
nrfipboot=$(cat "/sys/block/${pkfipboot}/$(basename ${fipbootdevice})/partition" 2>/dev/null)
[ -z "${nrfipboot}" ] && exit 1
partlabelfipboot=$(blkid ${fipbootdevice} -s PARTLABEL -o value)
[ -z "${partlabelfipboot}" ] && exit 1
typefipboot=$(echo ${partlabelfipboot} | cut -d "-" -f3)
[ -z "${typefipboot}" ] && typefipboot="${partlabelfipboot}"
echo "Type of bootchain installed (fip or boot): ${typefipboot}"

#########################################
# Convert boot to FIP or BOOT partition #
#########################################

if [ -n "$fip2boot" ] || [ -n "$boot2fip" ]; then
  rm -rf "${tmp}/boot"; mkdir "${tmp}/boot"
  echo "Copying files from /boot to ${tmp}/boot"
  cp -rfT "/boot/" "${tmp}/boot"
  sync
  while mountpoint -q "/boot"; do umount -R "/boot"; sleep 0.1; done
  if [ -n "$fip2boot" ]; then
    parted -s -- "/dev/${pkfipboot}" name ${nrfipboot} boot \
                                      set ${nrfipboot} boot on
    parted -s -- "/dev/${pkrootdev}"  set ${nrrootdev} boot off
    sync
    mkfs.vfat -v -F 32 -S 512 -s 16 -n "${target^^}-BOOT" "${fipbootdevice}"
    rm -rf "/boot"; mkdir "/boot"; chmod 0700 "/boot"
    sync
    mountoptions="rw,nosuid,nodev,noexec,relatime,nosymfollow,fmask=0077,dmask=0077,codepage=437"
    mountoptions+=",iocharset=ascii,shortname=mixed,utf8,errors=remount-ro"
    while mount "${fipbootdevice}" /boot -o "${mountoptions}" 2>/dev/null; ! mountpoint -q /boot
    do echo "Waiting for /boot being mounted..."; sleep 0.1; done
    typefipboot="boot"
  else
    parted -s -- "/dev/${pkfipboot}" name ${nrfipboot} fip \
                                      set ${nrfipboot} boot off
    parted -s -- "/dev/${pkrootdev}"  set ${nrrootdev} boot on
    rm -rf /boot; mkdir /boot; chmod 0700 /boot
    typefipboot="fip"
  fi
  sync
  fix_partlabelfipboot
  sync
  echo "Copying files from ${tmp}/boot to /boot"
  cp -rfT "${tmp}/boot" "/boot"
  sync
fi

##########################################
# Write linux or u-boot to fip partition #
##########################################
if [ -n "$write2fip" ] && [[ "${pkgbase}" == "${linuxpkg}" ]]; then
  if [[ "${typefipboot}" == "fip" ]]; then
    fipfiles=""
    [ -f "${bl31file}" ] && fipfiles+=" --soc-fw ${bl31file}"
    if [ -f "/boot/u-boot.bin" ]; then 
      fipfiles+=" --nt-fw /boot/u-boot.bin"
    elif [ -f "${linux}" ]; then
      fipfiles+=" --nt-fw ${linux}"
      if [[ "$(imageid ${linux})" == "5241644d" ]]; then # We have a linux kernel image
        [ -f "${atfdtb}" ] && fipfiles+=" --nt-fw-config ${atfdtb}"
        [ -f "${initrd}" ] && fipfiles+=" --tos-fw-extra2 ${initrd}"
      fi
    fi
    fiptool --verbose create "${tmp}/fip.bin" ${fipfiles}
    fiptool info "${tmp}/fip.bin"
    echo "Writing FIP to: ${fipbootdevice}"
    dd bs=4M of="${fipbootdevice}" if="${tmp}/fip.bin" conv=fsync
  fi
fi

########################
# Write ATF bootloader #
########################
if [ -n "$write2atf" ]; then
  if [ -n "${atfdevice}" ]; then
    headerdev="/dev/"$(lsblk -no pkname ${atfdevice})
    if [ -f "${headerfile}" ]; then
      echo "Writing ${headerfile} to ${headerdev}"
      dd if="${headerfile}" of="${headerdev}" conv=fsync
    fi
    echo -e "Target = ${target}, ATF device = ${device}"
    if [ -f "${atffile}" ]; then
      echo "Zeroing: ${atfdevice}"
      dd bs=64k if=/dev/zero of="${atfdevice}" conv=fsync 2>/dev/null
      echo "Writing ${atffile} to ${atfdevice}"
      dd bs=64k if="${atffile}" of="${atfdevice}" conv=fsync
    else
      echo "Atf binary does not exist: ${atffile}"
    fi
    fix_emmcbootpart "${headerdev}"
    fix_partlabelfipboot
  fi
fi

sync
[ -n "$create_parts" ] && parted -s -- "/dev/${pkrootdev}" print

exit 0
