#!/bin/sh
# Part of SMLinux Linux Distribution
# http://git.pktsurf.in/imagebuilder 

# Copyright (c) 2022-23 PktSurf <smlinux@pktsurf.in>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
############

# Abort on any error
set -e

basedir="$(pwd)"
tmp="/tmp"
smversion="1.0"
kernelversion="5.4.41"
build=1sml
arch="$( uname -m )"

#### DONT CHANGE THESE THREE UNLESS YOU WANT TO RENDER YOUR SYSTEM INOPERABLE!! ####
block0="/dev/loop0" 	# block0 is the loop0 device
block1="/dev/loop0p1"	# block1 is to be the FAT partition
block2="/dev/loop0p2"	# block2 is to be the EXT4 partition
####################################################################################

# Whether to create a highly-compressed zip file, yes/no
createzip="no"

# name the .img file
distroimg="smlinux-$smversion-$arch.img"

# Echo what we are about to do
echo "Creating '$arch' bootable image..."
echo ""

sleep 1

# Check whether we have the x86_64 kernel, mbr.bin in $basedir and /bin/extlinux
if [[ $arch = "x86_64" ]]; then
  echo "Checking whether we have the kernel..."
  if [[ ! -f $basedir/kernel-$kernelversion-$arch-$build.t?z ]]; then
    echo "$arch kernel not found. Aborting!"
    exit 1
  else
    echo "Found Kernel $kernelversion installer"
  fi

  echo "Checking for mbr.bin in $basedir..."
  if [[ ! -f $basedir/extlinux/mbr.bin ]]; then
    echo "mbr.bin not found. Aborting!"
    exit 1
 else
    echo "Found mbr.bin found"
  fi
  
  echo "Checking for the Extlinux bootloader binary..."
  if [[ ! -x /bin/extlinux ]]; then
    echo "Extlinux bootloader binary not found. Aborting!"
    exit 1
  else
    echo "Found Extlinux bootloader binary..."
  fi
	
  echo "Checking for extlinux.conf..."
  if [ ! -f $basedir/extlinux/extlinux.conf ]; then
    echo "extlinux.conf not found. Aborting!"
    exit 1
  else
    echo "extlinux.conf found, continuing..."
  fi
fi


# Uncompressed general size of the image (1G/2G/3G)
imgsize="2G"

# Ascertain whether we have an sfdisk partition table for the above size
if [ -f $basedir/sfdisk.part."$imgsize.$arch" ]; then
  echo "Found sfdisk.part."$imgsize.$arch", proceeding with image creation..."
else
  echo "sfdisk.part."$imgsize.$arch" was not found."
  echo "Create one using fallocate manually, run losetup, partition it and then:"
  echo "sfdisk -d $block0 > sfdisk.part.$imgsize.$arch"
  echo "Aborting!"
  exit 1
fi

# Just to be on the safe side, delete any previous .img files
echo "Removing any old .img and .zip files..."
rm -fv "$tmp"/"$distroimg"
rm -fv "$tmp"/"$distroimg".zip

# just a precautionary check, we don't intend to find an existing 
# $block1. If we do, we abort, cos someone's currently working on it
# or was, and didn't remove it.

if [[ -b $block1 ]]; then
  echo "$block1 was found. Someone has been working on it previously. Kindly remove it."
  echo "ABORTING!"
  exit 1
fi

# create a unique .img file
# the .img file is to be created in $TMP
# pass on the partition size to fallocate 
echo "Creating blank image using fallocate..."
fallocate -l $imgsize "$tmp/$distroimg" 

if [ "$?" != "0" ] ; then
  echo "Raw image creation failed. Aborting!"
fi

# Flush the buffers to storage
sync

# create the loop device
echo "Creating the loop0 device using losetup"
losetup -P -f "$tmp/$distroimg" 

# create temporary mount points (directories)
rootfstemp="$(mktemp -d $tmp/rootfs.XXXXX)"

# for partitioning, we redirect an already created partition file to sfdisk
# 100M for FAT, 900M for ext4/LINUX as a crude example (FOR ARM!)
# note: the output was derived after manually partitioning using cfdisk, 
# and then doing sfdisk -d /dev/loop0 > sfdisk.part.$arch
# be sure to remove the /dev/loop0* at the top of the file, second line
# and /dev/loop0* before "start=" words from that file.
sfdisk "$block0" < $basedir/sfdisk.part."$imgsize.$arch"
if [ "$?" = "0" ]; then
  echo "Partition table successfully written to the loop device"
else
  echo "Attempt to write the partition table to the loop device failed!"
  echo "Aborting!"
  exit 1
fi

sleep 1

# For the raspberry pi, we have created two partitions
# partition 1 /dev/loop0p1 to be formatted with FAT32,
# partition 2 /dev/loop0p2 to be formatted with ext4
#
# For x86_64, we have only *ONE* partition
# partition 1 /dev/loop0p1 to be formatted with only ext4

# -b is for block device checking
# format $block1 using the FAT file system first
if [[ -b $block1 ]] ; then
  echo "Found "$block1", formatting it..."
  if [[ $arch = "aarch64" ]]; then
    mkfs.fat "$block1" 
    if [[ $? = "0" ]] ; then
      echo "Formatting "$block1" with FAT FS successful"
    else
      echo "Formatting "$block1" with FAT FS FAILED!"
      exit 1
    fi
  elif [[ $arch = "x86_64" ]] ; then
    mkfs.ext4 "$block1"
    if  [[ $? = "0" ]] ; then
      echo "Formatting "$block1" with EXT4 FS succesful"
    else
      echo "Formatting "$block1" with EXT4 FS FAILED!"
      exit 1
    fi
  fi

else
  echo "$block1 not found. Is the sfdisk partition table correct?"
  exit 1
fi

# next, format $block2 as ext4
# assume this is a raspberry pi image only
if [[ -b $block2 ]]; then
  if [[ $arch = "aarch64" ]]; then
    mkfs.ext4 "$block2"
    if [[ $? = "0" ]]; then
      echo "Formatting "$block2" with EXT4 FS successful"
    else
      echo "Formatting "$block2" with EXT4 FS FAILED!"
      exit 1
    fi
  fi
fi

# next, we mount $block1/$block2 to their respective unique directories inside $TMP
# first mount the rootfs, then make a boot directory inside it if there isn't one, then
# mount the FAT partition on top of rootfs's boot directory as "boot".
# Again, $block1 is FAT FS, $block2 is EXT4 FS!
# Also, bear in mind that there is no block2 for x86_64 arch!

if [[ $arch = "aarch64" ]]; then
  mount -v "$block2" "$rootfstemp" 
  mkdir -p "$rootfstemp"/boot
  mount -v "$block1" "$rootfstemp"/boot/ 

elif [[ $arch = "x86_64" ]]; then
  mount -v "$block1" "$rootfstemp" 
fi

# Sort out the x86_64 stuff first. Install the bootloader, then the kernel. Other packages can wait.
if [[ $arch = "x86_64" ]]; then
  echo "Copying bootloader files..."
  mkdir -p "$rootfstemp"/boot
  cp -v "$basedir"/extlinux/{libutil.c32,menu.c32,extlinux.conf,initrd.img} "$rootfstemp"/boot/
  echo "Next, installing the bootloader..."
  # $block0 is /dev/loop0
  dd if=$basedir/extlinux/mbr.bin of="$block0"
  # $block1 is /dev/loop0p1
  extlinux -i "$rootfstemp"/boot
  echo "Installing the $arch kernel first, initfs will be copied later on..."
  installpkg --root "$rootfstemp" "$basedir"/kernel-$kernelversion-$arch-$build.t?z
	
  # UUID stuff
  # grab /dev/loop0p1's UUID. 
  uuidr="$( blkid -o value -s UUID $block1)"
  echo -n "$block1 s UUID is $uuidr , appending it to extlinux.conf"
  echo "APPEND root=UUID=$uuidr ro quiet consoleblank=0" >> "$rootfstemp"/boot/extlinux.conf
  echo "Done installing the bootloader."
fi

# do stuff here, like installing other packages, both, for AARCH64 and x86_64
# initfs shall be the first package.
installpkg --root "$rootfstemp" "$basedir/initfs-0.1-noarch-?sml.t?z" 
installpkg --root "$rootfstemp" "$basedir/packages/$arch/*/*" 
installpkg --root "$rootfstemp" "$basedir/packages/$arch/*/*/*" 

# Sync and sleep a little
sync
sleep 0.5

# Put any other miscellaneous stuff you want put in here
#cp -v "$basedir/packages/$arch/a/gawk-4.2.1-aarch64-1sml.txz" "$rootfstemp"/root/
#cp -v "$basedir/mkshrc" "$rootfs"/root/.mkshrc

# unmount them all when the job has been done
# specifically unmount "$rootfstemp"/boot in the case of ARM
# "$rootfstemp" is going to be get unmounted anyways
if [[ $arch = "aarch64" ]]; then
  umount -v "$rootfstemp"/boot 
fi
sleep 2
umount -v "$rootfstemp" 

# to debug the FS, comment this
rm -rfv "$rootfstemp"

# delete the loop device
losetup -d $block0

echo 
echo "$tmp/$distroimg successfully created!"

if [[ $CREATEZIP = "yes" ]]; then
  echo -n "Creating highly compressed zip file..."
  cd $tmp
  zip -9 "$distroimg".zip "$distroimg"
  echo "done"
fi

# done