#!/usr/bin/bash
#
# SPDX-FileCopyrightText: 2022 Celeste Liu <CoelacanthusHex@gmail.com>
# SPDX-License-Identifier: GPL-3.0-or-later

# shellcheck disable=1091
. /usr/share/makepkg/util.sh
colorize

verbose=0
use_fixed_password=0
build_firmware=1
varbose_arg=
rootfs="parchriscv-$(date --rfc-3339=date).tar.zst"
fstype=ext4
kernel=linux
kernel_suffix=""
cmdline=""

uboot_version=v2023.10
opensbi_version=v1.3.1

show_help() {
cat << EOF
Usage: ${0##*/} [-hvfd] [-p PASSWORD] [-r ROOTFS] [-t FSTYPE] [FILENAME]
Create Parch RISC-V distro image.

    FILENAME    generated image file name
                default: 'parchriscv-$(date --rfc-3339=date).qcow2'
                unless the extension is qcow2, implies raw disk format

    -h          display this help and exit
    -f          use fixed password instead of using systemd-firstboot to ask
    -d          only build the disk image and omit building OpenSBI/U-Boot
    -p PASSWORD set root password to PASSWORD instead of passwd in rootfs
    -r ROOTFS   specify rootfs file name
    -t FSTYPE   specify rootfs filesystem type (default: ext4)
    -k KERNEL   specify kernel package name (default: linux)
    -c CMDLINE  append CMDLINE to kernel command line
    -v          verbose mode
EOF
}

parse-args() {
    local OPTIND=1
    while getopts 'hvfdr:t:p:k:c:' opt; do
      case $opt in
        h)  show_help
            exit 0
            ;;
        v)  verbose=$((verbose+1))
            varbose_arg="--verbose"
            ;;
        f)  use_fixed_password=1
            ;;
        d)  build_firmware=0
            ;;
        p)  password=$OPTARG
            ;;
        r)  rootfs=$OPTARG
            ;;
        t)  fstype=$OPTARG
            ;;
        k)  kernel=$OPTARG
            ;;
        c)  cmdline=$OPTARG
            ;;
        *)  show_help >&2
            exit 1
            ;;
      esac
  done
  shift "$(( OPTIND-1 ))"
  filename=${1:-parchriscv-$(date --rfc-3339=date).qcow2}
}

toggle-systemd-firstboot() {
    msg2 "Toggle systemd-firstboot..."
    sudo rm -f mnt/etc/{machine-id,hostname,shadow}
    # Mask Arch filesystem's factory to avoid the removed shadow file being copied over
    sudo ln -s /dev/null mnt/etc/tmpfiles.d/arch.conf
}

use-fixed-password() {
    msg2 "Using fixed password... $password"
    [[ -n $password ]] && sudo usermod --root $(realpath ./mnt) --password $(openssl passwd -6 "$password") root
}

parse-args "$@"

if [ $build_firmware == 1 ]
then
    msg "Building U-Boot..."

    if [[ -d u-boot ]]; then
        pushd u-boot
        git checkout HEAD -- ':(top)'
        git fetch origin
        git checkout $uboot_version
        popd
    else
        git clone --depth 1 --filter=blob:none -b $uboot_version https://github.com/u-boot/u-boot.git
    fi
    pushd u-boot

    make \
        CROSS_COMPILE=riscv64-linux-gnu- \
        qemu-riscv64_smode_defconfig
    ./scripts/config \
        -e CMD_BTRFS -e FS_BTRFS
    make \
        CROSS_COMPILE=riscv64-linux-gnu- \
        olddefconfig
    make CROSS_COMPILE=riscv64-linux-gnu-
    popd

    msg "Building OpenSBI..."

    if [[ -d opensbi ]]; then
        pushd opensbi
        git checkout HEAD -- ':(top)'
        git fetch origin
        git checkout $opensbi_version
        popd
    else
        git clone --depth 1 --filter=blob:none -b $opensbi_version https://github.com/riscv-software-src/opensbi
    fi
    pushd opensbi

    # If user set it, OpenSBI will failed to build.
    unset PYTHONSAFEPATH

    make \
        CROSS_COMPILE=riscv64-linux-gnu- \
        PLATFORM=generic \
        FW_PAYLOAD_PATH=../u-boot/u-boot.bin
    popd

    cp ./opensbi/build/platform/generic/firmware/fw_payload.bin opensbi_fw_payload.bin
fi

if [[ $filename == *.qcow2 ]]
then
    msg "Create QCOW2 image..."
    qemu-img create -f qcow2 "$filename" 10G
    sudo modprobe nbd max_part=16 || exit 1
    # Possible NBD device collision?
    sudo qemu-nbd -c /dev/nbd0 "$filename"
    loopdev=/dev/nbd0
else
    msg "Create raw image..."
    fallocate "$filename" -l 10G
    loopdev=$(sudo losetup --show -P -f "$filename")
fi

msg "Partitioning..."
# FIXME: otherwise NBD device is not ready
sleep 1s
sudo parted "$loopdev" mklabel gpt mkpart "" "$fstype" 0% 100%

sudo partprobe "$loopdev"
partdev="$loopdev"p1

sudo mkfs.$fstype -L rootfs "$partdev"

sudo mkdir -p mnt
sudo mount "$partdev" mnt
sudo chown root:root mnt

uuid=$(sudo findmnt mnt -o UUID -n)

msg "Extract rootfs..."

pushd mnt
sudo bsdtar $varbose_arg -kpxf "../$rootfs"
popd

msg "Install kernel package..."

if [[ $kernel != linux ]]
then
    kernel_suffix="${kernel#linux}"
    echo -e "[unsupported]\nInclude = /etc/pacman.d/mirrorlist" | sudo tee -a mnt/etc/pacman.conf
fi

sudo systemd-nspawn -D mnt pacman \
    --noconfirm --needed \
    -Syu $kernel linux-firmware

sudo mkdir -p mnt/boot/extlinux
cat << EOF | sudo tee mnt/boot/extlinux/extlinux.conf
menu title Parch RISC-V Boot Menu
timeout 100
default linux-fallback

label linux
    menu label Linux linux
    kernel /boot/vmlinuz-linux$kernel_suffix
    initrd /boot/initramfs-linux$kernel_suffix.img
    append earlyprintk rw root=UUID=$uuid rootwait console=ttyS0,115200 $cmdline

label linux-fallback
    menu label Linux linux (fallback initramfs)
    kernel /boot/vmlinuz-linux$kernel_suffix
    initrd /boot/initramfs-linux$kernel_suffix-fallback.img
    append earlyprintk rw root=UUID=$uuid rootwait console=ttyS0,115200 $cmdline
EOF

cat << EOF | sudo tee mnt/etc/systemd/network/default.network
[Match]
Name=en*

[Network]
DHCP=yes
EOF
sudo systemd-nspawn -D mnt systemctl enable systemd-networkd.service

msg "Clean up..."
msg2 "Clean up pacman package cache..."
yes y | sudo pacman \
    --sysroot ./mnt \
    --sync --clean --clean

(( use_fixed_password==0 )) && toggle-systemd-firstboot || use-fixed-password

msg2 "Unmount and sync..."
sudo umount mnt
if [[ $filename == *.qcow2 ]]
then
    sudo qemu-nbd -d "$loopdev"
else
    sudo losetup -d "$loopdev"
fi
sudo sync