#!/bin/bash

set -e

[ "$(id -u)" = "0" ] ||
	{ echo "sorry, must be root"; exit 1; }

cleanup() {
	[ ! -d "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}
rq() {
   local out="${TEMP_D}/out"
	"$@" > "$out" 2>&1 || { error "FAILED [$?]:" "$@"; cat "$out"; return 1; }
}
fail() { msg "$@" 1>&2; exit 1; }
error() { echo "$@" 1>&2; }
msg() {
    local up idle
    read  up idle </proc/uptime
    #logger --tag=test-mic "[$up] $*"
    error "[$up] $*"
}

do_pt_test() {
    ## do_pt_test(img_in, pt_num, data-content, info-content, prefix)
    ##  - mic partition $pt and write data-content to a file 'data.txt'
    ##  - mic again and read file verifying it is there.
	local img="$1" pt="$2" data="$3" premsg="$4" out=""
	msg "${premsg}testing partition $pt in $img"
	mount-image-callback --cd "--part=$pt" "$img" -- \
		sh -c "echo '$data' > data.txt" || {
			error "${premsg}failed writing to partition $pt in $img";
			return 1;
		}
	local ret="" expected=$(printf "%s\n" "$data")
	msg "${premsg}wrote data to $pt"
	out=$(mount-image-callback --read-only --cd "--part=$pt" "$img" -- \
		sh -c '
		    r=0; for f in "$@"; do
		    cat $f || { echo "$f: cat failed rc=$?"; r=99; }; done;
		    exit $r' sh-extract data.txt)
	ret=$?
	if [ $ret -ne 0 -a $ret -ne 99 ]; then
		error "${premsg}failed mounting part $pt for verification";
		error "out=${out}"
		return 1;
	fi
	[ "$expected" = "$out" ] || {
		error "${premsg}"
		error "== expected on $pt =="
		error "$expected"
		error "== found on $pt =="
		error "$out"
		return 1
	}
	return 0
}

runloop() {
	local prefix="$1" img_in="$2" num=$3 i=0
	local img="${TEMP_D}/work.img"
	cp "$img_in" "${img}"
    # repeatedly write to partition 1 and partition 2 and verify.
    while [ $i -lt $num ]; do
        i=$((i+1))
	    do_pt_test "$img" 1 "foo1 $i" "$i $prefix: " ||
		    fail "${prefix}: failed testing partition 1 on $img"
	    do_pt_test "$img" 2 "foo2 $i" "$i $prefix: " ||
		    fail "${prefix}: failed testing partition 2 on $img"
    done
    msg "tested $prefix in $img successfully for $num times"
	return 0
}

TEMP_D=$(mktemp -d ${TMPDIR:-/tmp}/${0##*/}.XXXXXX)
trap cleanup EXIT

pt1="${TEMP_D}/pt1.img"
pt2="${TEMP_D}/pt2.img"
prept="${TEMP_D}/header.img"
postpt="${TEMP_D}/foot.img"
img_mbr="${TEMP_D}/disk-mbr.img"
img_gpt="${TEMP_D}/disk-gpt.img"
MB=$((1024*1024))
SSIZE=512
pt1_size=$((100*$MB))
pt2_size=$((200*$MB))
prept_size=$MB
postpt_size=$MB

pt1_d="${TEMP_D}/pt1"
pt2_d="${TEMP_D}/pt2"
mkdir -p "$pt1_d" "$pt2_d"
echo "partition 1" > "$pt1_d/info.txt"
echo "file 1" > "$pt1_d/file1.txt"
mkdir "$pt1_d/dev" "$pt1_d/sys" "$pt1_d/proc" "$pt1_d/mnt"
echo "partition 2" > "$pt2_d/info.txt"

## Stage 1
## Create 2 un-partitioned images (one mbr, one gpt) put a filesystem on them.
truncate "--size=$pt1_size" "$pt1"
truncate "--size=$pt2_size" "$pt2"
rq mkfs.ext4 -F "${pt1}"
rq mkfs.ext4 -F "${pt2}"
## Stage 2
## Create a full disk image with those 2 partition images inside
## and a partition table that points to them.  Do one for MBR and GPT.
truncate "--size=$prept_size" "$prept"
truncate "--size=$postpt_size" "$postpt"

msg "writing hunks to disk image ${img_mbr}"
for hunk in "$prept" "$pt1" "$pt2" "$postpt"; do
	rq dd bs=1M conv=notrunc oflag=append "if=$hunk" "of=${img_mbr}" ||
		fail "failed adding $hunk to disk image"
done
cp "${img_mbr}" "$img_gpt"

#
sfdisk_in="$TEMP_D/ptable_in"
curstart=0
(
echo unit: sectors
for pair in "-:${prept_size}" "1:${pt1_size}" "2:${pt2_size}" "-:${postpt_size}"; do
	op=${pair%%:*}
	size="${pair#*:}"
	if [ "$op" != "-" ]; then
		echo "${curstart} $((size/$SSIZE)) L"
	fi
	curstart=$(($curstart+($size/$SSIZE)))
done
) > "$sfdisk_in"

msg "partitioning MBR disk image ${img_mbr}"
(echo "label: dos"; cat "$sfdisk_in"; ) |
   rq sfdisk "${img_mbr}"

msg "partitioning GPT disk image ${img_gpt}"
(echo "label: gpt"; cat "$sfdisk_in"; ) |
   rq sfdisk "${img_gpt}"

## Stage 2.5
## Verify we can mount each partition with '--part=N'
## write data, then remount and read expected data.
runloop MBR "${img_mbr}" 100
#runloop GPT "${img_gpt}" 100

# vi: ts=4 noexpandtab