#!/bin/sh
# - nr_ssd
# - nr_hdd
# - nr_brd
# - nr_srd
# - nr_pmem
# - wait_disks_timeout

# the hosts config may either
# - hard code    xxx_partitions; leave nr_xxx_partitions empty
# - hard code nr_xxx_partitions; make     xxx_partitions a shell glob to be expanded in test box

[ -n "$wait_disks_timeout" ] || wait_disks_timeout=30

. $LKP_SRC/lib/unit.sh
. $LKP_SRC/lib/debug.sh

check_nr_partitions()
{
	max_ssd=$(ls $ssd_partitions | wc -l)
	max_hdd=$(ls $hdd_partitions | wc -l)

	nr_ssd_partitions=${nr_ssd_partitions:-$(echo $ssd_partitions | wc -w)}
	nr_hdd_partitions=${nr_hdd_partitions:-$(echo $hdd_partitions | wc -w)}

	[ "$ssd_partitions" ] && [ "$nr_ssd_partitions" != "$max_ssd" ] && {
		die "disk number $nr_ssd_partitions mismatch with $max_ssd real disks: $ssd_partitions"
	}

	[ "$hdd_partitions" ] && [ "$nr_hdd_partitions" != "$max_hdd" ] && {
		die "disk number $nr_hdd_partitions mismatch with $max_hdd real disks: $hdd_partitions"
	}

	[ "$nr_hdd_partitions" -gt 0 ] && [ -z "$hdd_partitions" ] && {
		die "Invalid hdd disk configuration, please check host file!"
	}

	[ "$nr_ssd_partitions" -gt 0 ] && [ -z "$ssd_partitions" ] && {
		die "Invalid ssd disk configuration, please check host file!"
	}

	return 0
}

wait_disks()
{
	for i in $(seq $wait_disks_timeout)
	do
		check_nr_partitions 2>/dev/null && return
		sleep 1
	done

	check_nr_partitions && return
	exit 1
}

check_nr_pmem()
{
	max_pmem=$(ls /sys/block/pmem* | wc -w)
	[ $max_pmem -eq 0 ] && die "No pmem device"

	[ "$nr_pmem" -gt "$max_pmem" ] && nr_pmem=$max_pmem
}

partition2disk()
{
	local part=$1
	local pat=

	case "$part" in
		/dev/nvme*) pat='p[0-9]+';;
		/dev/disk/by-id/*) pat='-part[0-9]+';;
		/dev/dm-*|/dev/mapper/*);;
		# General fallback pattern with /dev/sda1 /dev/loop0, etc.
		*) pat='[0-9]+';;
	esac
	if [ -n "$pat" ]; then
		echo $part | sed -re "s/$pat$//"
	else
		echo $part
	fi
}

[ -z "$nr_hdd" ] && nr_hdd=0
[ -z "$nr_ssd" ] && nr_ssd=0
[ -z "$nr_pmem" ] && nr_pmem=0

[ $((nr_hdd + nr_ssd)) -gt 0 ] && {
	# TODO: wait ssd and hdd on demand
	wait_disks
	[ "$nr_ssd" -gt "$max_ssd" ] && nr_ssd=$max_ssd
	[ "$nr_hdd" -gt "$max_hdd" ] && nr_hdd=$max_hdd
}

if [ -n "$nr_hdd" -a "$nr_hdd" != 0 ]; then
	nr_partitions=$nr_hdd
	partitions="$hdd_partitions"
	disk_description=${nr_hdd}hdd
elif [ -n "$nr_ssd" -a "$nr_ssd" != 0 ]; then
	nr_partitions=$nr_ssd
	partitions="$ssd_partitions"
	disk_description=${nr_ssd}ssd
elif [ -n "$nr_brd" -a "$nr_brd" != 0 ]; then
	[ -n "$ramdisk_size" ] || ramdisk_size="1G"
	log_cmd modprobe -r brd
	log_cmd modprobe brd rd_nr=$nr_brd rd_size=$(to_kb $ramdisk_size)
	sleep 1
	partitions=$(ls /dev/ram*)
	nr_partitions=$nr_brd
	disk_description=${nr_brd}brd
	disks=$(echo $partitions | tr '\n' ' ')
elif [ -n "$nr_srd" -a "$nr_srd" != 0 ]; then
	[ -n "$ramdisk_size" ] || ramdisk_size="1G"
	log_cmd modprobe -r scsi_debug
	log_cmd modprobe scsi_debug add_host=1 num_tgts=$nr_srd max_luns=1 dev_size_mb=$(to_mb $ramdisk_size)
	sleep 1
	partitions=$(grep -l -e scsi_debug /sys/block/*/device/model | awk -F'/' '{print "/dev/"$4}')
	nr_partitions=$nr_srd
	disk_description=${nr_srd}srd
	disks=$(echo $partitions | tr '\n' ' ')
elif [ -n "$scsi_ram" -a "$scsi_ram" != 0 ]; then
	[ -n "$ramdisk_size" ] || ramdisk_size="1G"
	log_cmd modprobe -r scsi_ram
	log_cmd modprobe scsi_ram capacity=$(($(to_byte $ramdisk_size) / 512))
	sleep 1
	partitions=$(grep -l -e "RAM Drive" /sys/block/*/device/model | awk -F'/' '{print "/dev/"$4}')
	nr_partitions=1
	disk_description=1scsi_ram
	disks=$(echo $partitions | tr '\n' ' ')
elif [ -n "$nr_pmem" -a "$nr_pmem" != 0 ]; then
	log_cmd modprobe nd_e820
	sleep 1
	check_nr_pmem
	nr_partitions=$nr_pmem
	partitions=$(ls /dev/pmem*)
	disk_descriptions=${nr_pmem}pmem
	disks="fake pmem disk, you shouldn't use it"
else
	echo "setup/disk: dumping shell variables:"
	set
	echo

	die "setup/disk: should not reach here"
fi

[ -z "$nr_partitions" -o "$nr_partitions" -lt 1 ] && {
	die "setup/disk: no disk partitions found, nr_partitions=$nr_partitions"
}

[ "${partitions#/dev/disk/by-}" != "$partitions" ] && partitions="$(readlink -e $partitions)"

partitions=$(echo $partitions | cut -f1-$nr_partitions -d' ')

if [ -z "$disks" ]; then
	for part in $partitions; do
		ret_disk=$(partition2disk $part)
		[ "${disks#*$ret_disk }" = "$disks" ] && disks="$disks $ret_disk "
	done
fi

cat >> $TMP/env.yaml <<EOF

# setup/disk
nr_partitions: $nr_partitions
partitions: $partitions
disk_description: $disk_description
disks: $disks
EOF
