#!/bin/bash
### rootfsinstall: create partitions and fileysytems
#
# Based on rpi-install by: Davide Madrisan
# Modified by: Silvan Calarco <silvan.calarco@mambasoft.it>
#

SYSTEMFSTYPE=ext4
# All targets are arm currently
TARGET_ARCH=arm

function usage() {
   echo "Usage: $0 device rootfs-archive [-bpi|-cubox|-rpi] [-a archive] [-c command] [-fs fstype] [-k kernel] [-p password] [-y]"
   echo
   echo "Target platform (optional):"
   echo "-bpi        : make a BananaPi SD with boot partition"
   echo "-cubox      : make a Cubox SD with boot partition"
   echo "-rpi        : make a RaspberryPi SD with boot partition"
   echo
   echo "Options:"
   echo "-a archive  : additional archive to unpack in rootfs root"
   echo "-cp command : command to run before chroot (may be specified multiple times)"
   echo "              Use @ROOT@ as reference to rootfs path and @BOOT@ as reference to bootfs path"
   echo "-c command  : command to run through chroot in rootfs (may be specified multiple times)"
   echo "-fs fstype  : filesystem type for system partition (default: ext4)"
   echo "-h hostname : set given hostname"
   echo "-k kernel   : extract and use specified kernel file from ROOTFS/boot/ for boot"
   echo "-k2 kernel  : extract and use specified second kernel file from ROOTFS/boot/ for boot"
   echo "-i packages : install packages using package manager in chroot"
   echo "-l          : force local installation only; don't perform installation from network"
   echo "-p password : set given password for root user"
   echo "-r mode     : set given default video mode where applicable (e.g. -r 1920x1080-32@60)"
   echo "-y          : assume yes to all questions (script mode)"
   echo
   echo "Example: $0 /dev/sdb /tmp/openmamba-rootfs-light.arm.tar.xz"
   exit 1
}

[ $UID -eq 0 ] || {
   echo "ERROR: this script must be called as root; aborting."
   exit 1
}

while [ "$1" ]; do
   case $1 in
      -y) ASSUMEYES=1 ;;
      -cubox) TARGETDEVICE=cubox ;;
      -bpi) TARGETDEVICE=bpi ;;
      -rpi) TARGETDEVICE=rpi ;;
      -fs) [ "$2" ] || {
              echo "Error: -fs option requires an argument; aborting."
              exit 1
           }
           SYSTEMFSTYPE=$2
           shift
           ;;
      -a) [ "$2" ] || {
              echo "Error: -a option requires an archive as argument; aborting."
              exit 1
           }
           ADDITIONALARCHIVE=`realpath $2`
           [ -r "$ADDITIONALARCHIVE" ] || {
              echo "Error: file $ADDITIONALARCHIVE not found; aborting."
              exit 1
           }
           shift
           ;;
      -c) [ "$2" ] || {
              echo "Error: -c option requires a command string as argument; aborting."
              exit 1
           }
           CHROOTCOMMANDS[${#CHROOTCOMMANDS[*]}]="$2"
           shift
           ;;
      -cp) [ "$2" ] || {
              echo "Error: -cp option requires a command string as argument; aborting."
              exit 1
           }
           PRECHROOTCOMMANDS[${#PRECHROOTCOMMANDS[*]}]="$2"
           shift
           ;;
      -k) [ "$2" ] || {
              echo "Error: -k option requires a kernel file name as argument; aborting."
              exit 1
           }
           USEKERNEL=$2
           shift
           ;;
      -k2) [ "$2" ] || {
              echo "Error: -k2 option requires a kernel file name as argument; aborting."
              exit 1
           }
           USEKERNEL2=$2
           shift
           ;;
      -h) [ "$2" ] || {
              echo "Error: -h option requires a hostname as argument; aborting."
              exit 1
           }
           SETHOSTNAME=$2
           shift
           ;;
      -i) [ "$2" ] || {
              echo "Error: -i option requires a packages list as argument; aborting."
              exit 1
           }
           INSTALLPACKAGES="$2"
           shift
           ;;
      -l) LOCALONLY=1 ;;
      -p) [ "$2" ] || {
              echo "Error: -p option requires a password as argument; aborting."
              exit 1
           }
           ROOTPASSWORD=$2
           shift
           ;;
      -r) [ "$2" ] || {
              echo "Error: -r option requires a video mode as argument; aborting."
              exit 1
           }
           VIDEOMODE=$2
           shift
           ;;
      -*) echo "Error: invalid option $1; aborting."
          exit 1
          ;;
       *) if [ ! "$DRIVE" ]; then
             DRIVE=$1
          elif [ ! "$ROOTFSARCHIVE" ]; then
             ROOTFSARCHIVE=`realpath $1`
          else
             echo "Error: invalid argument $1; aborting."
             exit 1
          fi
          ;;
    esac
    shift
done
[ "$ROOTFSARCHIVE" ] || { usage; exit 1; }

if [ "${DRIVE:0:11}" = "/dev/mmcblk" -o "${DRIVE:0:9}" = "/dev/loop" ]; then
   PARTITION1=${DRIVE}p1
   PARTITION2=${DRIVE}p2
else
   PARTITION1=${DRIVE}1
   PARTITION2=${DRIVE}2
fi

case $TARGETDEVICE in
   bpi) ARCH=arm
        BOOTFSTYPE=fat32
        BOOTFSSIZE=22M
        BOOTFSSTART=1049k
        [ "$USEKERNEL" ] || USEKERNEL=uImage-sunxi
        KERNELDEST="@BOOT@/kernel.img"
        INSTALLPACKAGES="kernel-sunxi kernel-sunxi-modules u-boot-sunxi-Bananapi $INSTALLPACKAGES"
        FSTAB_ROOT_DEVICE="/dev/mmcblk0p2"
        ;;
   rpi) ARCH=arm
        BOOTFSTYPE=fat32
        BOOTFSSIZE=40M
        BOOTFSSTART=0
        [ "$USEKERNEL" ] || USEKERNEL="zImage-*-rpi"
        [ "$USEKERNEL2" ] || USEKERNEL2="zImage-*-rpi-v7"
        KERNELDEST="@BOOT@/kernel.img"
        KERNELDEST2="@BOOT@/kernel7.img"
        INSTALLPACKAGES="kernel-rpi kernel-rpi-modules raspberrypi-firmware \
kernel-rpi-v7 kernel-rpi-v7-modules \
raspberrypi-utils raspi-config $INSTALLPACKAGES"
        RPI_FIRMWARE_FILES="bootcode.bin fixup_x.dat start_x.elf \
bcm27*.dtb cmdline.txt config.txt"
        FSTAB_ROOT_DEVICE="/dev/mmcblk0p2"
        ;;
   cubox) ARCH=arm
          BOOTFSTYPE=ext3
          BOOTFSSIZE=20M
          BOOTFSSTART=0
          [ "$USEKERNEL" ] || USEKERNEL=uImage-3.6.9-cubox
          KERNELDEST="@BOOT@/uImage"
          INSTALLPACKAGES="kernel-cubox kernel-cubox-modules xf86-video-dove $INSTALLPACKAGES"
          ;;
   *) BOOTFSTYPE=
      if [ "$USEKERNEL" ]; then
         KERNELDEST="/tmp/$USEKERNEL"
      else
         KERNELDEST=
      fi
      PARTITION2=$PARTITION1 ;;
esac

[ -r "$ROOTFSARCHIVE" ] || {
   echo "ERROR: unable to read rootfs archive $ROOTFSARCHIVE; aborting."
   exit 1
}

echo "Rootfs archive : $ROOTFSARCHIVE"
echo "Target SD card : ${DRIVE}"
[ "$TARGETDEVICE" ] && echo "Target platform: $TARGETDEVICE"
[ "$BOOTFSTYPE" ] &&   echo "Boot fs        : $BOOTFSTYPE($BOOTFSSIZE)"
[ "$BOOTFSTYPE" ] &&   echo "System fs      : $SYSTEMFSTYPE"
[ "$USEKERNEL" ] &&    echo "Kernel         : $USEKERNEL -> $KERNELDEST"
[ "$USEKERNEL2" ] &&   echo "Second kernel  : $USEKERNEL2 -> $KERNELDEST2"
echo

[ "$ASSUMEYES" ] || {
   echo -n "All current data in ${DRIVE} will be lost, do you want to proceed [y/N]?"
   read ans
   [ "$ans" = "y" -o "$ans" = "Y" ] || exit 0
}

sudo parted -s ${DRIVE} mktable msdos
if [ "$?" != "0" ]; then
   echo "ERROR: parted exited with error $? while creating partition table on ${DRIVE}; aborting."
fi
sync

if [ "$BOOTFSTYPE" ]; then
   # FAT partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $BOOTFSTYPE $BOOTFSSTART $BOOTFSSIZE
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating boot partition on ${DRIVE}; aborting."
   fi
   sudo parted -s -- ${DRIVE} toggle 1 boot
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while setting boot flag on ${DRIVE} boot partition; aborting."
   fi
   # rootfs partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $SYSTEMFSTYPE $BOOTFSSIZE -1
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating system partition on ${DRIVE}; aborting."
   fi
else
   # rootfs partition
   sudo parted -s -a cyl -- ${DRIVE} mkpart primary $SYSTEMFSTYPE 0 -1
   if [ "$?" != "0" ]; then
      echo "ERROR: parted exited with error $? while creating system partition on ${DRIVE}; aborting."
   fi
fi

if [ "$BOOTFSTYPE" ]; then
   echo "Formatting $PARTITION1 ($BOOTFSTYPE) ..."
   case "$BOOTFSTYPE" in
      vfat|fat32) mkfs.vfat -s 1 -F 32 -n "boot" $PARTITION1 || {
               echo "ERROR: unable to format $PARTITION1; aborting."
               exit 1
            } ;;
      ext*) mke2fs -q -t $BOOTFSTYPE -j -L "boot" $PARTITION1 || {
               echo "ERROR: unable to format $PARTITION1; aborting."
               exit 1
            } ;;
   esac
fi

echo "Formatting $PARTITION2 ($SYSTEMFSTYPE)..."
mke2fs -q -F -t $SYSTEMFSTYPE -j -L "openmamba" $PARTITION2 || {
   echo "ERROR: unable to format $PARTITION2; aborting."
   exit 1
}

[ "$BOOTFSTYPE" ] && BOOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`
ROOTMP=`mktemp -d -t rootfsinstall.XXXXXXXX`

[ -d $ROOTMP ] || {
   echo "Error: unable to create temporary mount points; aborting."
   exit 1
}

mount $PARTITION2 $ROOTMP || {
   echo "ERROR: unable to mount partition $PARTITION2; aborting."
   exit 1
}
trap "[ "$ROOTMP" -a -e $ROOTMP ] && umount $ROOTMP/dev/pts $ROOTMP/dev $ROOTMP/proc $ROOTMP/sys $ROOTMP/run $ROOTMP && rmdir $ROOTMP;[ "$BOOTMP" -a -e $BOOTMP ] && umount $BOOTMP && rmdir $BOOTMP" 0 HUP INT QUIT ABRT KILL TERM

if [ "$BOOTFSTYPE" ]; then
   mount $PARTITION1 $BOOTMP || {
      echo "ERROR: unable to mount partition $PARTITION1; aborting."
      exit 1
   }
fi

(
   cd $ROOTMP >/dev/null || {
      echo "ERROR: unable to access system partition mount point; aborting."
      umount $PARTITION2
      rmdir $ROOTMP
      exit 1
   }

   echo "Uncompressing $ROOTFSARCHIVE..."
   echo -n "This can take a while, please wait..."
   tar xf $ROOTFSARCHIVE || {
      echo
      echo "ERROR: unable to uncompress rootfs archive; aborting."
      exit 1
   }
   echo "DONE"
   for b in dev dev/pts proc sys run; do
      mount -o bind /$b $ROOTMP/$b
   done

   if [ "$ADDITIONALARCHIVE" ]; then
      echo "Uncompressing $ADDITIONALARCHIVE..."
      echo -n "This can take a while, please wait..."
      tar xf $ADDITIONALARCHIVE || {
         echo
         echo "ERROR: unable to uncompress additional archive; aborting."
         exit 1
      } 
      echo "DONE"
   fi
)

[ $? -eq 0 ] || exit $?

[ -e /usr/bin/qemu-${TARGET_ARCH}.static ] && cp /usr/bin/qemu-${TARGET_ARCH}.static $ROOTMP/usr/bin/qemu-${TARGET_ARCH}

if [ "$ROOTPASSWORD" ]; then
   echo "$ROOTPASSWORD" | chroot $ROOTMP passwd root --stdin || {
      echo "ERROR: unable to set password for root; aborting."
      exit 1
   }
fi

if [ "$SETHOSTNAME" ]; then
   echo "$SETHOSTNAME" > $ROOTMP/etc/hostname
   sed -i "s|HOSTNAME=.*|HOSTNAME=$SETHOSTNAME|" $ROOTMP/etc/sysconfig/network
   sed -i "s|\(127\.0\.0\.1\W*\)\(.*\)|\1 $SETHOSTNAME.localdomain $SETHOSTNAME localhost.localdomain localhost|" $ROOTMP/etc/hosts
fi

if [ "$INSTALLPACKAGES" -a "$LOCALONLY" != "1" ]; then
   mv $ROOTMP/etc/resolv.conf $ROOTMP/etc/resolv.conf.rootfsinstall
   cp /etc/resolv.conf $ROOTMP/etc/resolv.conf
   chroot $ROOTMP bash -c /etc/profile.d/gpg-key-openmamba.sh
   #chroot $ROOTMP smart channel --enable devel-kernel
   chroot $ROOTMP smart update || {
      echo "ERROR: error running chroot command 'smart update'; aborting."
      exit 1
   }
   chroot $ROOTMP smart install -y $INSTALLPACKAGES || {
      echo "ERROR: error running chroot command 'smart install $INSTALLPACKAGES'; aborting."
      exit 1
   }
   #chroot $ROOTMP smart channel --disable devel-kernel
   mv $ROOTMP/etc/resolv.conf.rootfsinstall $ROOTMP/etc/resolv.conf
fi

KERNELDEST=`echo $KERNELDEST | sed "s|@BOOT@|$BOOTMP|g"`
KERNELDEST2=`echo $KERNELDEST2 | sed "s|@BOOT@|$BOOTMP|g"`

#
# Prepare boot partition
#
if [ "$TARGETDEVICE" = "bpi" ]; then
   # TODO
   cp $ROOTMP/boot/Bananapi-script.bin $BOOTMP/
   IMAGE=`ls $ROOTMP/boot/uImage-*-sunxi | head -n 1`
   cp $IMAGE $BOOTMP/uImage-sunxi
   cat > $BOOTMP/uEnv.txt << _EOF
bootargs=console=ttyS0,115200 console=tty0 sunxi_g2d_mem_reserve=0 sunxi_ve_mem_reserve=0 disp.screen0_output_mode=EDID:1280x720p50 hdmi.audio=EDID:0 console=tty1 root=/dev/mmcblk0p2 rootfstype=ext4 elevator=deadline rootwait
aload_script=fatload mmc 0 0x43000000 Bananapi-script.bin;
aload_kernel=fatload mmc 0 0x48000000 uImage-sunxi; bootm 0x48000000;
uenvcmd=run aload_script aload_kernel
_EOF
   dd if=$ROOTMP/boot/Bananapi/u-boot-sunxi-with-spl.bin of=${DRIVE} bs=1024 seek=8
   :
elif [ "$TARGETDEVICE" = "rpi" ]; then
   # fetch and install firmware files if missing
   for f in $RPI_FIRMWARE_FILES; do
      if [ -e $ROOTMP/boot/$f ]; then
         cp $ROOTMP/boot/$f $BOOTMP/$f
      elif [ "$f" != "cmdline.txt" -a "$f" != "config.txt" ]; then
         echo "ERROR: firmware file $f for rpi is missing in /boot directory; aborting."
         exit 1
      fi
   done
   if [ -e $ROOTMP/boot/overlays ]; then
      cp -a $ROOTMP/boot/overlays $BOOTMP/
   fi
#      IMAGE=`ls $ROOTMP/boot/$USEKERNEL | head -n 1`
   cp $ROOTMP/boot/$USEKERNEL $KERNELDEST
   cp $ROOTMP/boot/$USEKERNEL2 $KERNELDEST2

   # cmdline.txt : parameters passed to the kernel on boot
   [ ! -e $ROOTMP/boot/cmdline.txt ] && \
   echo "\
dwc_otg.lpm_enable=0 \
root=/dev/mmcblk0p2 \
rootfstype=$SYSTEMFSTYPE \
quiet \
rootwait \
splash" > $BOOTMP/cmdline.txt

   # Create the kernel configuration file
   [ ! -e $ROOTMP/boot/config.txt ] && \
   echo "\
start_file=start_x.elf
fixup_file=fixup_x.dat
framebuffer_depth=32
framebuffer_ignore_alpha=1
[rp1]
arm_freq=800
force_turbo=1
gpu_mem=100
disable_overscan=1
[rp2]
gpu_mem=128
" > $BOOTMP/config.txt

   ## Enable HDMI mode rather than DVI, if required
   #echo "\
   #hdmi_drive=2" >> $BOOTMP/config.txt

   # optional files:
   #   config.txt    : A configuration file read by the GPU.
   #                   Use this to override set the video mode, alter system clock
   #                   speeds, voltages, etc.
   #   vlls directory: Additional GPU code, e.g. extra codecs. 
   #                   Not present in the initial release.

   [ "$SWHOME" -a "${SWHOME:0:4}" = "/tmp" ] && rm -rf $SWHOME

elif [ "$TARGETDEVICE" = "cubox" ]; then
   if [ -e $ROOTMP/boot/uImage-*-cubox ]; then
      IMAGE=`ls $ROOTMP/boot/uImage-*-cubox | head -n 1`
      cp $IMAGE* $BOOTMP/
      ln -s ${IMAGE/*\/}-boot.scr $BOOTMP/boot.scr
   else
      [ "$VIDEO_MODE" ] || VIDEO_MODE="1920x1080-32@60"
      mkdir -p $BOOTMP
      echo "Creating u-boot boot.scr..."
      cat > $BOOTMP/boot.txt << _EOF
setenv bootargs 'console=ttyS0,115200n8 vmalloc=384M root=/dev/mmcblk0p2 video=dovefb:lcd0:$VIDEO_MODE-edid clcd.lcd0_enable=1 clcd.lcd1_enable=0 rootwait'
echo ======== Loading kernel ========
ext2load mmc 0:1 0x00200000 /uImage
echo ======== Booting kernel ========
bootm
_EOF
      mkimage -A arm -O linux -T script -C none -n "u-boot commands" -d $BOOTMP/boot.txt $BOOTMP/boot.scr || {
         echo "ERROR: mkimage: error creating boot.scr; aborting."
         exit 1
      }
   fi
fi

if [ "$FSTAB_ROOT_DEVICE" ]; then
   cat >> $ROOTMP/etc/fstab << _EOF
$FSTAB_ROOT_DEVICE	/	auto	defaults	0	0
_EOF
fi

if [ "$USEKERNEL" ]; then
   if [ -r $ROOTMP/boot/$USEKERNEL ]; then
      cp $ROOTMP/boot/$USEKERNEL $KERNELDEST || {
         echo "ERROR: unable to find or copy kernel $USEKERNEL; aborting."
         exit 1
      }
      echo "Kernel image available as $KERNELDEST"
   else
      echo "WARNING: can't find kernel file $ROOTMP/boot/$USERKERNEL"
   fi
fi

#
# Final custom commands
#
for c in `seq 1 ${#PRECHROOTCOMMANDS[*]}`; do
   PRECHROOTCOMMAND=`echo ${PRECHROOTCOMMANDS[c-1]} | sed "s|@ROOT@|$ROOTMP|g;s|@BOOT@|$BOOTMP|g"`
   $PRECHROOTCOMMAND || {
      echo "ERROR: error running pre-chroot command '${PRECHROOTCOMMANDS[c-1]}'; aborting."
      exit 1
   }
done

for c in `seq 1 ${#CHROOTCOMMANDS[*]}`; do
   chroot $ROOTMP ${CHROOTCOMMANDS[c-1]} || {
      echo "ERROR: error running chroot command '${CHROOTCOMMANDS[c-1]}'; aborting."
      exit 1
   }
done

sync

echo "Done."
exit 0