#!/bin/bash
#
# Copyright (C) 2022-2024 Nethesis S.r.l.
# SPDX-License-Identifier: GPL-2.0-only
#

set -e

white=$'\033[0;97m'
green=$'\033[0;32m'

function print_help {
    echo "Usage: $0 [-t <device>] -f"
    echo "  -t <device>: install on the given block device"
    echo "  -f: force install, ignore existing data and override partitions"
    exit 1
}

function check {
    disk=$(basename $1)
    if [ $(grep -c $disk /proc/partitions) -gt 1 ]; then
        echo -e "Multiple partitions found on $1, all content on device will be deleted!"
        read -rp "Do you want to proceed? (y/N) " force
        case $force in
            [yY] ) return 0;;
            * ) echo "Aborted"; return 1;;
        esac
    fi
    return 0
}

function install() {
    # Find boot device and umount it
    part=$(lsblk -ln | grep "/rom" | awk '{print $1}')
    disk=$(lsblk -lno pkname /dev/$part)
    device="/dev/$disk"
    if [ $device == $1 ]; then echo "Error: can't install on source disk" >&2 ; exit 1; fi
    if mount | grep -q /boot; then
        umount /boot 2>/dev/null
        umount /boot 2>/dev/null
    fi
    umount /overlay
    umount /rom

    sleep 1
    count=$(parted -sm $device print 2>/dev/null | tail -n1 | cut -d: -f3 | sed "s/MB//")
    echo "Writing image... "
    dd if=$device of=$1 bs=1M count=$count
    # check if the installation disk partition has been modified
    if ! lsblk -n --raw $device | grep part | grep -q "^${disk}128"; then
        echo "Fixing partition table..."
        sync
        sleep 1
        sgdisk -r 3:128 $1
    fi
    sync
    echo 3 > /proc/sys/vm/drop_caches
    sleep 1
    echo "Installation complete. Halting... "
    halt
}


force=0
while getopts "hft::" opt; do
    case $opt in
        (f) force=1 ;; #Force write
        (t) target=${OPTARG} ;; #Target disk
        (h) print_help ;;
        (*) printf "Illegal option '-%s'\n" "$opt" && exit 1 ;;
    esac
done

# select existing disks
if [ -z "$target" ]; then
    echo "Select target disk (or press Enter to exit):"
    counter=1
    disks="$(lsblk -nlo NAME,SIZE,TYPE,MODEL | awk '$3 ~ /^disk/ && $1 !~ /ram/ && $2 !~ /^0B$/')"
    while read -r disk size type model; do
	used=$(grep -q $disk /proc/mounts && echo "used" || echo "free")
	echo -e "${green}${counter}) ${white}${disk}	$size	$used	$model"
	counter=$((counter+1))
    done <<< "$disks"

    read -rp "Your choice: " choice
    if [ ${choice:-0} -gt 0 ] && [ ${choice} -lt ${counter} ]; then
	target=$(echo "$disks" | awk 'NR=='$choice'{ print $1; }')
        target="/dev/${target}"
    else
       echo "Aborted: invalid choice"
       exit 1
    fi
fi

if [ -b "${target}" ]; then
    if [ "$force" -eq 1 ]; then
        install $target
    else
        check $target && install $target
   fi
else
    echo -e "Target device not found: '$target'"
    exit 1
fi
