#!/bin/bash

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

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 "${atfdevice}" ] && return
  local pkdev=$(lsblk -rno pkname "${atfdevice}")
  [ -z "${pkdev}" ] && return
  dev=$(blkid -t "PARTLABEL=$1" -o device $(parts "${pkdev}"))
  [ -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 write_extlinux {
  mkdir -p $(dirname "$1")
  cat <<-EOF | tee "$1"
	DEFAULT linux-bpir-git
	  MENU title U-Boot menu
	  PROMPT 0
	  TIMEOUT 50
	LABEL linux-bpir-git
	  MENU LABEL Archlinux ARM for BananaPi Routers
	  LINUX ${2}
	  INITRD ${3}
	  FDT ${4}
	EOF
}

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
}

#################
# Set variables #
#################
tmp="/tmp/bpir-toolbox-tmp"
rootdev=$(lsblk -pilno name,type,mountpoint | grep -G 'part /$' | head -n1 | cut -d " " -f1)
if [ -z "$rootdev" ]; then
  if [ -f "/etc/bpir-is-initrd" ]; then  # Writing NAND while running from initrd
    if [[ $@ =~ "--nand-" ]] || [[ $@ =~ "--download2root" ]]; then
      source /etc/bpir-is-initrd
      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 linux-${target}-git
      [[ $? == 1 ]] && download linux-bpir-git
      download bpir-atf-git
      download bpir-uboot-git
      bpir-initrd
    fi
  else
    echo "Need to mount '/' first (and '/boot' if applicable)"
    exit 1
  fi
else
  partlabelroot=$(blkid $rootdev -s PARTLABEL -o value)
  [ -z "$partlabelroot" ] && exit 1
  target=$(echo $partlabelroot | cut -d'-' -f1)
fi

case ${target} in
  bpir64) default_ddrsize=1;;
  bpir3)  default_ddrsize=2;;
  bpir3m) default_ddrsize=2;;
  bpir4)  default_ddrsize=4;;
  *)      echo "Unknown target ${target}"; exit 1;;
esac
default_dtb="${target/m/-mini}"
default_dtb="${default_dtb/bpir/-bananapi-bpi-r}"
default_dtb=$(shopt -s nullglob; cd /boot/dtbs; echo "mt"*"${default_dtb}.dtb" | head -n 1)
default_dtb="${default_dtb/.dtb/}"

[ ! -f "$(cat /boot/bootcfg/initrd 2>/dev/null)" ] && rm -f "/boot/bootcfg/initrd"

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

mkdir -p "${tmp}/dtbos"
mkdir -p /boot/bootcfg
dtb=$(     getconf r dtb     "$default_dtb")
device=$(  getconf r device  "$(echo $partlabelroot | cut -d'-' -f2)")
linux=$(   getconf w linux   "/boot/Image")
cmdline=$( getconf w cmdline "console=ttyS0,115200 debug=7 rw rootwait audit=0")
initrd=$(  getconf w initrd  "/boot/initramfs-bpir.img")
atfdtb=$(  getconf w atfdtb  "/boot/dtbs/${dtb}-atf.dtb")
ddrsize=$( getconf r ddrsize "$default_ddrsize")

########## cleanup *-atf.dtb

extra=""
[[ "$ddrsize" != "$default_ddrsize" ]] && 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"

atfdevice=$(getdev atf)
fipbootdevice=$(getdev fip)
[ -z "$fipbootdevice" ] && fipbootdevice=$(getdev boot)
if [ -n "$fipbootdevice" ]; then
  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 (fip or boot): $typefipboot"
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

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

if [[ $@ =~ "--fip2boot" ]] || [[ $@ =~ "--boot2fip" ]]; then
  rm -rf $tmp/boot; mkdir $tmp/boot
  echo "Copying files from /boot to $tmp/boot"
  cp -rfT /boot/ $tmp/boot
  sync
  [ -z "$fipbootdevice" ] && exit 1
  while mountpoint -q /boot; do umount -R /boot; sleep 0.1; done
  if [[ $@ =~ "--fip2boot" ]]; then
    parted -s -- "/dev/$pkfipboot" name $nrfipboot boot \
                               set  $nrfipboot boot on
    sync
    mkfs.vfat -v -F 32 -S 512 -s 16 -n "${target^^}-BOOT" ${fipbootdevice}
    rm -rf /boot; mkdir /boot; chmod 0700 /boot
    sync
    while mount ${fipbootdevice} /boot 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
    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 ATF bootloader #
########################
if [[ $@ =~ "--atf" ]]; 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

###############
# Format NAND #
###############
if [[ $@ =~ "--nand-format" ]]; then
  [[ "$(tr -d '\0' 2>/dev/null </proc/device-tree/compatible)" != "bananapi,"* ]] && exit 1
  i=0
  ubinr=$(mtdnr ubi)
  ubidetach -p /dev/mtd${ubinr}
  ubiformat -y /dev/mtd${ubinr}
  ubiattach -p /dev/mtd${ubinr}
  ubidev=$(get_ubidev ${ubinr})
  [[ $? != 0 ]] && exit 1
  [[ "$ubidev" == "/dev/ubi_ctrl" ]] && exit 1
  ubimkvol ${ubidev} -n $i -N fip       -s 1MiB   -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
  while [ ! -e "${ubidev}_$i" ]; do sleep 0.1; done
  mkfs.ubifs ${ubidev}_$i
fi


###############
# Update NAND #
###############
if [[ $@ =~ "--nand-update" ]] || [[ $@ =~ "--nand-format" ]]; then
  [[ "$(tr -d '\0' 2>/dev/null </proc/device-tree/compatible)" != "bananapi,"* ]] && exit 1
  ubinr=$(mtdnr ubi)
  bl2nr=$(mtdnr bl2)
  atfbin="/usr/share/bpir-atf/${target}-atf-spim-nand-atf${extra}.bin"
  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
  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)
  fipfiles=""
  nandbl31file="/usr/share/bpir-atf/${target}-atf-spim-nand-bl31${extra}.bin"
  [ -f "$nandbl31file" ] && fipfiles+=" --soc-fw $nandbl31file"
  fipfiles+=" --nt-fw /usr/share/bpir-uboot/u-boot-${target}-emmc.bin"
  fiptool --verbose create $tmp/fip.bin $fipfiles
  fiptool info $tmp/fip.bin
  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
  ubirootfs=$(ubivol rootfs)
  mkdir -p $tmp/mnt
  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
  [ -f "${linux}.gz" ] && ubootlinux="${linux}.gz" || ubootlinux="${linux}"
  write_dtb "$tmp/atf.dtb" \
            "root=" \
            "/boot/dtbs/${dtb}.dtb" \
	    /usr/share/buildR64arch/boot/${target^^}/*.dts \
	    /usr/share/buildR64arch/boot/${target^^}-NAND/*.dts
  ubiupdate "$tmp/atf.dtb" "$tmp/mnt${atfdtb}"
  write_extlinux "$tmp/extlinux.conf" \
                 "${ubootlinux}" \
                 "${initrd}" \
                 "${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

############
# UARTBOOT #
############
if [[ $@ =~ "--uartboot" ]]; then
  mkdir -p /tmp/uartboot
  rm -rf /tmp/uartboot/*
  write_dtb "$tmp/atf.dtb" \
            "root=" \
            "/boot/dtbs/${dtb}.dtb" \
	    /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"
  #echo "Creating Image.xz..."
  #xz -e9cC crc32 ${linux} > $tmp/Image.xz
  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
fi

############
# EXTLINUX #
############
if [[ $@ =~ "--extlinux" ]] || ( [[ $@ =~ "--uboot-install" ]] && [ ! -f "$extlinux" ] ) ; then
  if mountpoint -q /boot/; then skip=5; else skip=0; fi
  write_extlinux "${extlinux}" \
                 "${linux:$skip}" \
                 "${initrd:$skip}" \
                 "${atfdtb:$skip}"
  [[ $@ =~ "--extlinux" ]] && exit 0
fi

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

if [[ $@ =~ "--uboot-remove" ]]; then
  rm -vf "/boot/u-boot.bin"
  rm -vf $extlinux
  rmdir -v $(dirname $extlinux)
fi

#########################################
# Create single dtb and write fip image #
#########################################
if [[ $@ =~ "--write2fip" ]] || [[ $@ =~ "--boot2fip" ]] || [[ $@ =~ "--fip2boot" ]]; then
  rm -v /boot/dtbos/*.dtbo /boot/dtbs/*-fixed.dtb 2>/dev/null # not used anymore
  write_dtb "$atfdtb" \
            "root=PARTLABEL=${partlabelroot}" \
            "/boot/dtbs/${dtb}.dtb" \
            /boot/dtbos/*.dts
  if [[ "${typefipboot}" == "fip" ]]; then
    fipfiles=""
    [ -f "$bl31file" ]   && fipfiles+=" --soc-fw $bl31file"
    [ -f "$linux" ]  && 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
    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

exit 0
