#!/bin/bash

set -euo pipefail

/preload-k8s-images.sh || exit 1

export CTL_BIN_LAYER=$( skopeo inspect docker-daemon:localhost/monok8s/control-agent:dev | jq -r '.Layers[0] | sub("^sha256:"; "")' )

mkdir -p \
  "$ROOTFS/dev" \
  "$ROOTFS/proc" \
  "$ROOTFS/sys" \
  "$ROOTFS/run" \
  "$ROOTFS/data" \
  "$ROOTFS/var" \
  "$ROOTFS/tmp" \
  "$ROOTFS/build" \
  "$ROOTFS/var/cache/apk" \
  "$ROOTFS/usr/lib/monok8s/crds" \
  "$ROOTFS/opt/monok8s/config"

mount --bind /var/cache/apk "$ROOTFS/var/cache/apk"
mount --bind /dev  "$ROOTFS/dev"
mount --bind /proc "$ROOTFS/proc"
mount --bind /sys  "$ROOTFS/sys"
mount --bind /run  "$ROOTFS/run"

cp /usr/bin/qemu-aarch64-static "$ROOTFS/usr/bin/"
cp /etc/resolv.conf "$ROOTFS/etc/resolv.conf"
cp /build/crio.tar.gz "$ROOTFS/build/"
cp /build/crds/*.yaml "$ROOTFS/usr/lib/monok8s/crds"

chroot "$ROOTFS" /bin/sh -c "ln -s /var/cache/apk /etc/apk/cache"
# chroot "$ROOTFS" /bin/sh -c "apk update"
chroot "$ROOTFS" /bin/sh -c "apk add bash curl"
cp "/install-packages.sh" "$ROOTFS/install-packages.sh"
chroot "$ROOTFS" /bin/bash /install-packages.sh || exit 1
rm "$ROOTFS/install-packages.sh"

/merge-rootfs.sh "/build/rootfs-extra" "$ROOTFS"

cp "/configure-system.sh" "$ROOTFS/configure-system.sh"
chroot "$ROOTFS" /bin/bash /configure-system.sh || exit 1
rm "$ROOTFS/configure-system.sh"

umount "$ROOTFS/var/cache/apk"
umount "$ROOTFS/dev"
umount "$ROOTFS/proc"
umount "$ROOTFS/sys"
umount "$ROOTFS/run"

rm -r "$ROOTFS/build"
rm "$ROOTFS/etc/resolv.conf"

### Begin making full disk image for the device
#!/bin/bash
set -euo pipefail

ROOTFS="${ROOTFS:?ROOTFS is required}"
BOARD_ITB="${BOARD_ITB:-/build/board.itb}"

IMG="${IMG:-output.img}"
DISK_SIZE="${DISK_SIZE:-8G}"

ROOTFS_IMG="${ROOTFS_IMG:-rootfs.ext4}"
ROOTFS_IMG_ZST="${ROOTFS_IMG_ZST:-rootfs.ext4.zst}"
ROOTFS_PART_SIZE_MIB="${ROOTFS_PART_SIZE_MIB:-2048}"

FAKE_DEV="/tmp/dev"
MNT_ROOTFS_IMG="/mnt/rootfs-img"
MNT_DATA="/mnt/data"

LOOP=""
ROOTFS_LOOP=""
TMP_LOOP=""
ROOTFS_TMP_LOOP=""

cleanup_fake_nodes() {
    local prefix="$1"
    [ -n "$prefix" ] || return 0
    find "$FAKE_DEV" -maxdepth 1 -type b -name "${prefix}*" -exec rm -f {} \; 2>/dev/null || true
}

cleanup() {
    set +e

    mountpoint -q "$MNT_ROOTFS_IMG" && umount "$MNT_ROOTFS_IMG"
    mountpoint -q "$MNT_DATA" && umount "$MNT_DATA"

    if [ -n "$ROOTFS_LOOP" ]; then
        losetup -d "$ROOTFS_LOOP" 2>/dev/null || true
    fi
    if [ -n "$LOOP" ]; then
        losetup -d "$LOOP" 2>/dev/null || true
    fi

    if [ -n "$ROOTFS_LOOP" ]; then
        cleanup_fake_nodes "$(basename "$ROOTFS_LOOP")"
    fi
    if [ -n "$LOOP" ]; then
        cleanup_fake_nodes "$(basename "$LOOP")"
    fi
}
trap cleanup EXIT

mkdir -p "$FAKE_DEV" "$MNT_ROOTFS_IMG" "$MNT_DATA"

echo "##################################################### Packaging RootFS $(du -sh "$ROOTFS" | awk '{print $1}')"

###############################################################################
# 1. Build reusable rootfs ext4 image once
###############################################################################

ROOTFS_BYTES=$(du -s -B1 "$ROOTFS" | awk '{print $1}')
EXTRA_BYTES=$((256 * 1024 * 1024))
IMG_BYTES=$(( ROOTFS_BYTES + ROOTFS_BYTES / 4 + EXTRA_BYTES ))

ALIGN=$((4 * 1024 * 1024))
IMG_BYTES=$(( (IMG_BYTES + ALIGN - 1) / ALIGN * ALIGN ))

MAX_BYTES=$(( ROOTFS_PART_SIZE_MIB * 1024 * 1024 ))
if [ "$IMG_BYTES" -ge "$MAX_BYTES" ]; then
    echo "ERROR: estimated rootfs image size $IMG_BYTES exceeds slot size $MAX_BYTES" >&2
    exit 1
fi

rm -f "$ROOTFS_IMG" "$ROOTFS_IMG_ZST"

truncate -s "$IMG_BYTES" "$ROOTFS_IMG"
mkfs.ext4 -F -L rootfs "$ROOTFS_IMG"

ROOTFS_LOOP=$(losetup --find --show -P "$ROOTFS_IMG")
/sync-loop.sh "$ROOTFS_LOOP"

# For a raw ext4 image there is usually no partition, so mount the loop device directly.
mount "$ROOTFS_LOOP" "$MNT_ROOTFS_IMG"

(
    cd "$ROOTFS"
    tar cpf - --exclude='./var' .
) | (
    cd "$MNT_ROOTFS_IMG"
    tar xpf -
)

mkdir -p "$MNT_ROOTFS_IMG/var"
mkdir -p "$MNT_ROOTFS_IMG/boot"
cp "$BOARD_ITB" "$MNT_ROOTFS_IMG/boot/kernel.itb"

sync
umount "$MNT_ROOTFS_IMG"

losetup -d "$ROOTFS_LOOP"
cleanup_fake_nodes "$(basename "$ROOTFS_LOOP")"
ROOTFS_LOOP=""

e2fsck -fy "$ROOTFS_IMG"
resize2fs -M "$ROOTFS_IMG"
e2fsck -fy "$ROOTFS_IMG"

echo "##################################################### Compressing OTA Image"
zstd -19 -T0 -f "$ROOTFS_IMG" -o "$ROOTFS_IMG_ZST"

tee >(wc -c > "${ROOTFS_IMG}.size") < "$ROOTFS_IMG" | sha256sum > "${ROOTFS_IMG}.sha256"
tee >(wc -c > "${ROOTFS_IMG_ZST}.size") < "$ROOTFS_IMG_ZST" | sha256sum > "${ROOTFS_IMG_ZST}.sha256"

cat > catalog.txt <<EOF
- version: $KUBE_VERSION
  url: "monok8s-$KUBE_VERSION-$TAG.ext4.zst"
  checksum: $(cut -d' ' -f1 < "${ROOTFS_IMG}.sha256")
  size: $(< "${ROOTFS_IMG}.size")
  zChecksum: $(cut -d' ' -f1 < "${ROOTFS_IMG_ZST}.sha256")
  zSize: $(< "${ROOTFS_IMG_ZST}.size")
  tag: "$TAG"
EOF

###############################################################################
# 2. Build full disk image
###############################################################################

rm -f "$IMG"
truncate -s "$DISK_SIZE" "$IMG"

sgdisk -o "$IMG" \
  -n 1:2048:+64M  -t 1:0700 -c 1:config \
  -n 2:0:+2G      -t 2:8300 -c 2:rootfsA \
  -n 3:0:+2G      -t 3:8300 -c 3:rootfsB \
  -n 4:0:0        -t 4:8300 -c 4:data

losetup -D
LOOP=$(losetup --find --show -P "$IMG")
/sync-loop.sh "$LOOP"

TMP_LOOP="$FAKE_DEV/$(basename "$LOOP")"

mkfs.vfat -F 32 -n MONOK8S_CFG "${TMP_LOOP}p1"
mkfs.ext4 -F -L rootfsB "${TMP_LOOP}p3"
mkfs.ext4 -F -L data "${TMP_LOOP}p4"

dd if="$ROOTFS_IMG" of="${TMP_LOOP}p2" bs=4M conv=fsync

# Grow each filesystem to fill its partition
e2fsck -fy "${TMP_LOOP}p2"
resize2fs "${TMP_LOOP}p2"

# Populate data partition
mount "${TMP_LOOP}p4" "$MNT_DATA"

cp -a "$ROOTFS/var" "$MNT_DATA/"
mkdir -p "$MNT_DATA/etc-overlay/work"
mkdir -p "$MNT_DATA/etc-overlay/upper"

sync
umount "$MNT_DATA"

losetup -d "$LOOP"
cleanup_fake_nodes "$(basename "$LOOP")"
LOOP=""

echo "Built artifacts:"
echo "  Full disk image: $IMG"
echo "  Rootfs OTA image: $ROOTFS_IMG"
echo "  Rootfs OTA compressed: $ROOTFS_IMG_ZST"

echo "##################################################### Compressing Full Disk Image"
gzip "/build/$IMG"