#!/bin/bash

rootdev=$(lsblk -pilno pkname,type,mountpoint | grep -G 'part /$' |  head -n1 | cut -d " " -f1)
partlabelroot=$(lsblk -plno partlabel $rootdev | grep -G '\-root$' )
partlabeluboot=$(lsblk -plno partlabel $rootdev | grep -G '\-uboot$' )
[ ! -n "$partlabeluboot" ] && exit
rkdev=${partlabeluboot#*-}
target=${partlabeluboot%"-$rkdev"}
rkdev=${rkdev/"-uboot"/""}
if [[ "$rkdev" =~ "@" ]]; then
  atfdevice=${rkdev#*@}
  rkdev=${rkdev%"@$atfdevice"}
fi

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

append="earlyprintk console=tty1 consoleblank=0 rw rootwait root=PARTLABEL=${partlabelroot}\
 audit=0 cgroup_enable=cpuset cgroup_memory=1 cgroup_enable=memory swapaccount=1"

case ${target} in
  rk3288)
    append+=' '""
    ;;
  rk3588)
    append+=' '"console=ttyS2,1500000n8 irqchip.gicv3_pseudo_nmi=0 switolb=1 coherent_pool=2M"
    ;;
esac

if [[ $@ =~ "--extlinux" ]] || [ ! -d $(dirname $extlinux) ]; then
  mkdir -p $(dirname $extlinux)
  thisarch=$(uname -m)
  if [[ $thisarch =~ "armv7" ]]; then
    default="linux-armv7"
  elif [[ $thisarch == "aarch64" ]]; then
    if pacman -Qs '^linux-aarch64-rk-rc$' > /dev/null ; then
      default="linux-aarch64-rk-rc"
    elif pacman -Qs '^linux-aarch64$' > /dev/null ; then
      default="linux-aarch64"
    else
      default="linux-aarch64-rk-rc"
      echo "WARNING: default linux package $default not found!"
    fi
  fi
  cat <<-EOF | tee $extlinux
	DEFAULT ${default}
	  MENU title U-Boot menu
	  PROMPT 0
	  TIMEOUT 50
	EOF
  if pacman -Qs '^linux-armv7$' > /dev/null ; then
    cat <<-EOF | tee -a $extlinux
	LABEL linux-armv7
	  MENU LABEL Archlinux armv7 linux
	  LINUX /zImage
	  FDT /dtbs/${target}-joined.dtb
	  APPEND ${append}
	  # Optionally add: video=1920x1080@60
	EOF
  fi
  if pacman -Qs '^linux-aarch64$' > /dev/null ; then
    cat <<-EOF | tee -a $extlinux
	LABEL linux-aarch64
	  MENU LABEL Archlinux aarch64 linux
	  LINUX Image
	  INITRD /initramfs-linux.img
	  FDT /dtbs/rockchip/${target}-${rkdev}-joined.dtb
	  APPEND ${append}
	  # Optionally add: video=1920x1080@60
	EOF
  fi
  for folder in $(shopt -s nullglob; echo /lib/modules/*) ; do
    pkgbase=$(cat $folder/pkgbase 2>/dev/null)
    if [ ! -z "$pkgbase" ]; then
      cat <<-EOF | tee -a $extlinux
	LABEL ${pkgbase}
	  MENU LABEL Archlinux ARM package ${pkgbase}
	  LINUX /vmlinuz-${pkgbase}
	  INITRD /initramfs-${pkgbase}.img
	  FDT /dtbs/${pkgbase}/rockchip/${target}-${rkdev}-joined.dtb
	  APPEND ${append}
	  # Optionally add: video=1920x1080@60
	EOF
    fi
  done
  exit
fi

if [[ $@ =~ "--install" ]]; then
  exit
fi

if [[ $@ =~ "--remove" ]]; then
  rm -vf /boot/dtbs/rk3288-openhour.dtb
  rm -vf /boot/dtbs/*/*/*-tojoin.dtb
  rm -vf /boot/dtbs/*/*/*-joined.dtb
  rm -vf /boot/dtbs/*/*-tojoin.dtb
  rm -vf /boot/dtbs/*/*-joined.dtb
  rm -vf /boot/dtbos
  exit
fi

filexz_sdmmc="/boot/uboot/u-boot-with-spl-${target}-${rkdev}.bin.xz"
filexz_emmc="/boot/uboot/u-boot-rockchip-${target}-${rkdev}.bin.xz"
filexz_spi="/boot/uboot/u-boot-rockchip-spi-${target}-${rkdev}.bin.xz"
[ ! -f "$filexz_sdmmc" ] && filexz_sdmmc=${filexz_sdmmc/"${target}-rock-"/"${target}-rock"}
[ ! -f "$filexz_emmc" ]  && filexz_emmc=${filexz_emmc/"${target}-rock-"/"${target}-rock"}
[ ! -f "$filexz_spi" ]   && filexz_spi=${filexz_spi/"${target}-rock-"/"${target}-rock"}

if [[ $@ =~ "--uboot@root" ]]; then
  ubootdev="/dev/disk/by-partlabel/${partlabeluboot}"
  [ ! -L "$ubootdev" ] && exit
  filexz="/doesnotexistanyway"
  [[ "${atfdevice}" == "sdmmc" ]] && filexz="$filexz_sdmmc"
  [[ "${atfdevice}" == "emmc" ]] && filexz="$filexz_emmc"
  if [[ "${atfdevice}" == "nvme" ]]; then
    echo "Write U-Boot to the used medium manually."
    echo "Use one of the following commands:"
    echo "  rockchip-toolbox --uboot@spi"
    echo "  rockchip-toolbox --uboot@emmc"
    echo "  rockchip-toolbox --uboot@sdmmc"
  fi
  [ ! -f "$filexz" ] && exit
  echo "Writing $(basename $filexz) to $(realpath ${ubootdev})"
  dd bs=64k of="${ubootdev}" if="/dev/zero" 2>/dev/null
  xz -dcv "$filexz" | dd bs=64k of="${ubootdev}"
  sync
fi

if [[ $@ =~ "--uboot@emmc" ]]; then
  [[ $(systemd-detect-virt) != "none" ]] && exit
  ubootdev="/dev/mmcblk0"
  [ ! -b "${ubootdev}boot0" ] && ubootdev="/dev/mmcblk1"
  [ ! -b "${ubootdev}boot0" ] && exit
  partlabeluboot=$(lsblk -plno partlabel $ubootdev | grep -G '\-uboot$' )
  ubootdev="/dev/disk/by-partlabel/${partlabeluboot}"
  [ ! -b "$ubootdev" ] && exit
  echo "Writing $(basename $filexz_emmc) to $(realpath ${ubootdev})"
  dd bs=64k of="${ubootdev}" if="/dev/zero" 2>/dev/null
  xz -dcv "$filexz_emmc" | dd bs=64k of="${ubootdev}"
  sync
fi

if [[ $@ =~ "--uboot@sdmmc" ]]; then
  [[ $(systemd-detect-virt) != "none" ]] && exit
  ubootdev="/dev/mmcblk0"
  [ -b "${ubootdev}boot0" ] && ubootdev="/dev/mmcblk1"
  [ -b "${ubootdev}boot0" ] && exit
  partlabeluboot=$(lsblk -plno partlabel $ubootdev | grep -G '\-uboot$' )
  ubootdev="/dev/disk/by-partlabel/${partlabeluboot}"
  [ ! -b "$ubootdev" ] && exit
  echo "Writing $(basename $filexz_sdmmc) to $(realpath ${ubootdev})"
  dd bs=64k of="${ubootdev}" if="/dev/zero" 2>/dev/null
  xz -dcv "$filexz_sdmmc" | dd bs=64k of="${ubootdev}"
  sync
fi

if [[ $@ =~ "--uboot@spi" ]]; then
  [[ $(systemd-detect-virt) != "none" ]] && exit
  ubootdev="/dev/mtdblock0"
  [ ! -b "$ubootdev" ] && exit
  if [ -f "$filexz_spi" ]; then
    echo "Writing $(basename $filexz_spi) to SPI. This may take a while..."
    xz -dcv "$filexz_spi" | dd of="${ubootdev}"
    sync
    exit
  else # Vendor U-boot uses same file as SD
    [ ! -f "$filexz_sdmmc" ] && exit
    echo "Writing $(basename $filexz_sdmmc) to SPI. This may take a while..."
    xz -dcv "$filexz_sdmmc" | dd of="${ubootdev}" bs=512 seek=64    count=1024 # Write SPL max 512 KiB
    xz -dcv "$filexz_sdmmc" | dd of="${ubootdev}" bs=512 seek=16384 skip=$((16384-64)) # Write U-Boot
    sync
    exit
  fi
fi

# The rest of the script for --dtbs argument
[[ $@ =~ "--dtbs" ]] || exit 0

# Create openhour dtb if possible
if [ -f "/boot/dtbs/rk3288-miqi.dtb" ]; then
  cp -vf /boot/dtbs/rk3288-miqi.dtb /boot/dtbs/rk3288-openhour.dtb
  echo "Editing rk3288-openhour.dtb for differences between Openhour Chameleon and Miqi"
  fdtput /boot/dtbs/rk3288-openhour.dtb -d   /mmc@ff0f0000 non-removable
  fdtput /boot/dtbs/rk3288-openhour.dtb -t u /mmc@ff0f0000 bus-width 4
  fdtput /boot/dtbs/rk3288-openhour.dtb -t s /mmc@ff0f0000 cap-sd-highspeed
  fdtput /boot/dtbs/rk3288-openhour.dtb -t x /mmc@ff0f0000 card-detect-delay 0xc8
fi

if [ ! -d "/boot/dtbos" ]; then
(
  mkdir -p "/boot/dtbos"
  cd /usr/share/buildRKarch/boot/dtbos-${target^^}/
  cp -vrf            ${target}-all-*.dts      "/boot/dtbos" 2>/dev/null
  cp -vrf --parent */${target}-all-*.dts      "/boot/dtbos" 2>/dev/null
  cp -vrf --parent   ${target}-${rkdev}-*.dts "/boot/dtbos" 2>/dev/null
  cp -vrf --parent */${target}-${rkdev}-*.dts "/boot/dtbos" 2>/dev/null
)
fi

# Create .dtbo's
rm -rf /boot/dtbos/*.dtbo /boot/dtbos/*/*.dtbo
for bp in $(shopt -s nullglob; echo /boot/dtbos/*.dts /boot/dtbos/*/*.dts); do
  echo "Creating .dtbo from $bp"
  dtc -q -I dts -O dtb -o ${bp/".dts"/".dtbo"} $bp
done

# Create joined .dtb
for file in $(find "/boot/dtbs" -name "${target}-${rkdev}.dtb"); do
  tojoin=${file/".dtb"/"-tojoin.dtb"}
  joined=${file/".dtb"/"-joined.dtb"}
  parts=(${file//\// })
  pkg="${parts[2]}"
  dtbos="$(shopt -s nullglob; echo /boot/dtbos/*.dtbo /boot/dtbos/${pkg}/*.dtbo)"
  if test -n "$dtbos"; then
    cp -vf $file $tojoin
    # Extra fdtput command inside .dts to be executed on the original .dtb
    for bp in $(shopt -s nullglob; echo /boot/dtbos/*.dts); do
      cat $bp | grep "//fdtput" | while read -r line ; do
        echo fdtput $tojoin ${line/"//fdtput"/""}
             fdtput $tojoin ${line/"//fdtput"/""}
      done
    done
    fdtoverlay -vi $tojoin -o $joined $dtbos
  else
    echo "No overlay applied!"
    cp -vf $file $joined
  fi

done

exit 0
