#!/bin/bash

function usage {
  cat <<-EOF
	Usage: $(basename "$0") [OPTION]...
	  -p --preset [PRESET]       specify preset
	  -P --allpresets            build all presets
	  -m --modulesonly           build only when image holds modules
	EOF
  exit 1
}

while [ "${1:-}" != '' ]; do
  case "$1" in
  '-p' | '--preset') shift; preset="${1}"      ;;
         '--preset'=*)      preset="${1#*=}"   ;;
  '-P' | '--allpresets')    allpresets=1       ;;
  '-m' | '--modulesonly')   modulesonly=1      ;;
  *) usage ;;
  esac
  shift
done

if [ -n "$allpresets" ]; then
  pkgbases=$(shopt -s nullglob; for f in /lib/modules/*/pkgbase ; do cat $f; echo; done | xargs)
elif [ -n "$preset" ]; then
  pkgbases="$preset"
else
  usage
fi

for pkgbase in $pkgbases; do

preset="/etc/mkinitcpio.d/${pkgbase}.preset"
if [ -f  "${preset}" ]; then
  source ${preset}
else
  default_config="/etc/mkinitcpio-bpir.conf"
  default_image="/boot/initramfs-${pkgbase}.img"
fi
if [ ! -f  "${default_config}" ]; then
  echo "Cannot open ${default_config}!"
  usage
fi
source ${default_config}
for hook in ${HOOKS[*]}; do
  [[ ! "${hook}" =~ "bpir" ]] && continue
  source "/etc/initcpio/hooks/${hook}"
done

binaries=(bash mount umount mountpoint ls cat dmesg date
  lsmod rmmod insmod modinfo modprobe depmod ldd xargs
  switch_root readlink basename dirname sync blkid setsid unshare
  cp mv dd mkdir mknod sed sleep ln rm uname touch chroot
  stat chown chmod debootstrap curl wget ar sha256sum pgrep
  grep head tail cut tee tr parted lsblk partprobe wipefs
  blkdiscard realpath gzip gunzip diff
  tar sort uniq wc zstdcat xzcat xz find mmc cpio zstd fiptool
  fsck.f2fs fsck.btrfs fsck.vfat
  mkfs.f2fs mkfs.btrfs mkfs.vfat mkfs.ubifs
  ubiattach ubidetach ubiupdatevol ubimkvol ubiformat ubinfo
  flash_erase lsusb lspci mkimage
  fdtget fdtput fdtoverlay ip bridge ethtool ping dhcpc
  bpir-toolbox bpir-initrd bpir-build bpir-rootfs
  nano '/usr/share/terminfo/l/linux'
  'libnss_files.so.2' 'libnss_dns.so.2' 'libresolv.so.2'
  '/usr/lib/debootstrap/pkgdetails'
)

for module in ${MODULES[*]}; do
  modfile=$(modinfo -k ${ALL_kver} -F filename ${module})
  [ -n "${modfile}" ] && modules+=("${modfile}")
done

if [[ ! ${modules[@]} ]] && [ -n $modulesonly ] && [ -f "$default_image" ] ; then
  echo "$(basename $0): No modules inside $(basename $default_image), skip rebuilding"
  exit 0
fi

wdir="/tmp/initramfs.$$"

function mkchroot
{
  local dest i ix d
  [ $# -lt 2 ] && return
  dest=$1
  shift
  for i in "$@"; do
    if [ "${i:0:1}" != "/" ]; then
      if [[ "$i" == "lib"*".so"* ]]; then
        i=$(find /lib/ -type f -name "$(basename "$i")" -print -quit)
      else
        i=$(command -v $i 2>/dev/null)
      fi
    fi
    [ -f "$dest/$i" ] && continue
    if [ -e "$i" ]; then
      mkdir -p "$dest/$(dirname "$i")" &&
      cat "$i" > "$dest/$i" &&
      [[ -x "$i" ]] && chmod +x "$dest/$i"
      mkchroot "$dest" $(ldd "$i" 2>/dev/null | grep -E -o '/.* ')
    elif [ -n "$i" ]; then
      ix=""; [ -e "${i}.zst" ] && ix="${i}.zst"; [ -e "${i}.xz"  ] && ix="${i}.xz"
      if [ -n "$ix" ]; then
        mkdir -p "$dest$(dirname ${ix})"
        zstd -d -q --output-dir-flat="$dest$(dirname ${ix})" "${ix}"
      else
        echo "Possibly missing binary: $i"
      fi
    fi
  done
}

echo "Building ${default_image}..."

rm -rf $wdir

mkdir -p      $wdir/{dev,etc,proc,run,sys,usr/lib,usr/bin,usr/local,usr/share,mnt/root,etc/modprobe.d,etc/rootcfg}
ln -s usr/bin $wdir/bin
ln -s usr/bin $wdir/sbin
ln -s ../bin  $wdir/usr/local/bin
ln -s ../bin  $wdir/usr/local/sbin
ln -s bin     $wdir/usr/sbin
ln -s usr/lib $wdir/lib
ln -s usr/lib $wdir/lib64
ln -s lib     $wdir/usr/lib64
mkchroot      $wdir ${binaries[*]} ${BINARIES[*]} ${FILES[*]} ${modules[*]}

ln -s /bin/bash $wdir/bin/sh

cp -f /etc/rootcfg/target  $wdir/etc/rootcfg/target
cp -f /etc/rootcfg/ddrsize $wdir/etc/rootcfg/ddrsize
echo -n "initrd"         > $wdir/etc/rootcfg/device
touch                      $wdir/etc/bpir-is-initrd

cp -rf /usr/share/debootstrap/ "${wdir}/usr/share"

if [[ "${MODULES_DECOMPRESS}" == "yes" ]]; then
  for modfile in ${modules[*]}; do
    [[ "${modfile}" != *".ko" ]] && zstd --rm -q -d "${wdir}${modfile}"
  done
fi

touch $wdir/etc/modprobe.d/modprobe.conf
echo "nameserver 8.8.8.8" > $wdir/etc/resolv.conf

echo -e -n '#!/bin/bash\n#' > "$wdir/init"
type run_hook >> "$wdir/init"
cat <<'EOT' | tee -a "$wdir/init" >/dev/null

PATH=/usr/bin
export PATH

mknod -m 640 /dev/console c 5 1
mknod -m 664 /dev/null    c 1 3
mount -n -t devtmpfs devtmpfs /dev
mount -n -t proc     proc     /proc
mount -n -t sysfs    sysfs    /sys
mount -n -t tmpfs    tmpfs    /run
ln -s /proc/self/fd /dev/fd

read -r cmdline < /proc/cmdline
init=/sbin/init
export root=
rootdelay=
rootfstype=auto
ro="ro"
rootflags=
device=
for param in $cmdline ; do
  case $param in
    init=*      ) init=${param#init=}             ;;
    root=*      ) export root=${param#root=}      ;;
    rootdelay=* ) rootdelay=${param#rootdelay=}   ;;
    rootfstype=*) rootfstype=${param#rootfstype=} ;;
    rootflags=* ) rootflags=${param#rootflags=}   ;;
    ro          ) ro="ro"                         ;;
    rw          ) ro="rw"                         ;;
  esac
done

run_hook

[ -n "$rootdelay" ] && sleep "$rootdelay"
mkdir /.root
[ -n "$rootflags" ] && rootflags+=","
rootflags+="$ro"
device="$root"
if [ -z "$device" ];then
  echo "No root device specified."
  setsid /bin/bash -c 'exec /bin/bash </dev/ttyS0 >/dev/ttyS0 2>&1'
fi
for (( i = 0 ; i < 50 ; i++ )); do
  [[ "$root" =~ "=" ]] && device=$(blkid -t $root -l -o device)
  [ -b "$device" ] && break || sleep 0.1
done

if ! mount -n -t "$rootfstype" -o "$rootflags" "$device" /.root ; then
  echo "Device $device could not be mounted!"
  setsid /bin/bash -c 'exec /bin/bash </dev/ttyS0 >/dev/ttyS0 2>&1'
else
  echo "Successfully mounted device $root"
fi

exec switch_root /.root "$init" "$@"
EOT
chmod +x $wdir/init

cat <<'EOT' | tee "$wdir/bin/reboot" >/dev/null
#!/bin/bash
/bin/sync
echo 1 > /proc/sys/kernel/sysrq
echo _sb > /proc/sysrq-trigger
EOT
chmod +x "$wdir/bin/reboot"

cat <<'EOT' | tee "$wdir/bin/bpir-synctime" >/dev/null
#!/bin/bash
date +"%d %b %Y %T %Z" -s "$(curl -s --head http://google.com | grep '^Date:' | cut -d' ' -f 3-)"
EOT
chmod +x "$wdir/bin/bpir-synctime"

cat <<'EOT' | tee "$wdir/bin/bpir-dhcpc" >/dev/null
#!/bin/bash

IFS=$'\n'

[ -z "$1" ] && intf="wan" || intf="$1"

echo "Setting $intf up and waiting for it to be up..."
ip link set $intf up
while [ -z "$(ip link show dev $intf up 2>/dev/null | \
	         grep 'state UP')" ]
do sleep 0.1; done

echo "Asking DHCP server for IP number."
outp=$(dhcpc -i $intf)
if [[ $? != 0 ]]; then
  echo "DHCP error on first call!"
  exit
fi
for line in $outp; do eval "${line}"; done

echo "Requesting lease of $YourIP."

outp=$(dhcpc -i $intf -r $YourIP)
if [[ $? != 0 ]]; then 
  echo "DHCP error on second call!"
  exit
fi
for line in $outp; do eval "${line}"; done

ip addr add $YourIP/$SubnetMask brd + dev $intf

ip route add default via ${DefaultGateways%%,*} dev $intf

echo "# initrd resolv.conf" >/etc/resolv.conf
for ns in ${DomainNameServers//,/ }; do
    echo "nameserver $ns" >>/etc/resolv.conf
done

while [ -z "$(ip addr show dev $intf 2>/dev/null | \
	         grep $YourIP)" ]
do sleep 0.1; done

/bin/bpir-synctime

ip addr show dev $intf
echo
ip route
echo
cat /etc/resolv.conf
EOT
chmod +x "$wdir/bin/bpir-dhcpc"

(cd $wdir; ls -AR)

echo -ne "\nCompressing ${default_image}, size "

[ -z "${COMPRESSION}" ] && COMPRESSION="gzip"
( cd $wdir; find . -print0 | cpio --null -o --format=newc | zstd --format=${COMPRESSION} -19 -T0) > $default_image

rm -rf $wdir

stat --printf="Done building ${default_image}, size %s bytes\n" $default_image

done
