#!/bin/sh
set -x
panic() {
    echo "panic: $@"
    exit 1
}

# Setup writable
if [ ! -d /writable ]; then
    mkdir /writable
    cp -R /system-data /writable/
    mkdir -p /writable/system-data/etc/writable/
    mkdir -p /writable/system-data/var/lib/cloud/seed/nocloud-net/
    cp /lxd/hostname /writable/system-data/etc/writable/
    cp /lxd/meta-data /lxd/network-config /lxd/user-data /lxd/vendor-data /writable/system-data/var/lib/cloud/seed/nocloud-net/
fi

# Grab the snaps version
SNAP_KERNEL=$(grep pc-kernel_ /system-data/var/lib/snapd/seed/seed.yaml | awk '{print $NF}')
SNAP_CORE=$(grep 'core[0-9]\{0,2\}_' /system-data/var/lib/snapd/seed/seed.yaml | awk '{print $NF}')

# Attempt to parse the grub environment
new_snap_mode=""
if [ -e  "/writable/system-data/boot/grub/grubenv" ]; then
    . /writable/system-data/boot/grub/grubenv
    [ -n "${snap_kernel}" ] && SNAP_KERNEL=${snap_kernel}
    [ -n "${snap_core}" ] && SNAP_CORE=${snap_core}

    if [ "${snap_mode}" = "try" ]; then
        new_snap_mode="trying"
        [ -n "${snap_try_kernel}" ] && SNAP_KERNEL=${snap_try_kernel}
        [ -n "${snap_try_core}" ] && SNAP_CORE=${snap_try_core}
    elif [ "${snap_mode}" = "trying" ]; then
        new_snap_mode=""
    fi
else
    mkdir -p /writable/system-data/boot/grub/
    printf "# GRUB Environment Block\n" > /writable/system-data/boot/grub/grubenv
    printf "#######################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################################" >> /writable/system-data/boot/grub/grubenv
fi

# Extract initramfs
mkdir -p /mnt /initrd
mount -t tmpfs tmpfs /initrd
/bin/squashfuse /writable/system-data/var/lib/snapd/snaps/${SNAP_KERNEL} /mnt -o ro,allow_other
cd /initrd
(lzcat | /bin/cpio -id) < /mnt/initrd.img >/dev/null 2>&1 || true
if [ ! -e "/initrd/scripts/" ]; then
    (/bin/cpio -t; lzcat | /bin/cpio -id) < /mnt/initrd.img >/dev/null 2>&1 || true
fi
if [ ! -e "/initrd/scripts/" ]; then
    (/bin/cpio -t; /bin/cpio -id) < /mnt/initrd.img >/dev/null 2>&1 || true
fi
if [ ! -e "/initrd/scripts/" ]; then
    panic "Unable to unpack initrd"
fi
umount /mnt
cd /

# Exec initramfs
[ ! -L /scripts ] && ln -s /initrd/scripts /scripts
. /initrd/scripts/ubuntu-core-rootfs
export PATH=$PATH:/initrd/bin:/initrd/sbin
export root=/root

# Setup the root
rootmnt=/root
rm -rf ${rootmnt}
mkdir -p ${rootmnt}

## Mount the root
/bin/squashfuse /writable/system-data/var/lib/snapd/snaps/${SNAP_CORE} /root -o ro,allow_other

## Mount runtime directory
mount -o rw,nosuid,noexec,relatime,mode=755 -t tmpfs tmpfs "/root/run"

## Mount the writable storage
mount --bind /writable /root/writable

## squashfuse workaround
if [ ! -d "/root/usr/local/bin" ]; then
    mount -t tmpfs tmpfs /root/usr/local -o size=100k,mode=0711
    mkdir -p /root/usr/local/bin/

    mkdir -p /root/usr/local/lib/
    for lib in /lib/libfuse* /lib/libsquashfuse*; do
        [ ! -e ${lib} ] && continue

        touch /root/usr/local/lib/$(basename ${lib})
        mount --bind ${lib} /root/usr/local/lib/$(basename ${lib})
    done
else
    mount -t tmpfs tmpfs /root/usr/local/bin/ -o size=100k,mode=0711
fi

touch /root/usr/local/bin/squashfuse
mount --bind /bin/squashfuse /root/usr/local/bin/squashfuse

touch /root/usr/local/bin/mount.fuse
mount --bind /bin/squashfuse /root/usr/local/bin/mount.fuse

touch /root/usr/local/bin/mount.real
mount --bind /root/bin/mount /root/usr/local/bin/mount.real

cat > /root/usr/local/bin/mount << EOF
#!/bin/sh
export LD_LIBRARY_PATH=/usr/local/lib/
if [ "\$1" = "-t" ]; then
    if echo \$@ | grep -q fuse.squashfuse || echo \$3 | grep -q '\.snap$'; then
        exec /usr/local/bin/mount.fuse \$3 \$4
    fi

    exec /usr/local/bin/mount.real \$@
fi

if echo \$@ | grep -q fuse.squashfuse || echo \$1 | grep -q '\.snap$'; then
    exec /usr/local/bin/mount.fuse \$1 \$2
fi

exec /usr/local/bin/mount.real \$@
EOF
chmod +x /root/usr/local/bin/mount
mount --bind /root/usr/local/bin/mount /root/bin/mount

## Prepare the fstab
fstab="${rootmnt}/etc/fstab"
writable_paths="${rootmnt}/etc/system-image/writable-paths"

## Add writable overlays
if [ -e "$writable_paths" ]; then
    touch "${rootmnt}/run/image.fstab"
    mount -o bind "${rootmnt}/run/image.fstab" "$fstab" || panic "Cannot bind mount fstab"
    echo "# Auto-generated by $0" >> "$fstab"
    echo "# DO NOT EDIT THIS FILE BY HAND - YOUR CHANGES WILL BE OVERWRITTEN" >> "$fstab"
    echo "# (See writable-paths(5) for details)" >> "$fstab"
    handle_writable_paths "$writable_paths" "$fstab"
fi

## Mount the systemd overlay so that we have a complete root partition during boot
mkdir -p "${rootmnt}/writable/system-data/etc/systemd/system"

if [ ! -e "${rootmnt}/writable/system-data/etc/systemd/system/systemd-hostnamed.service.d/override.conf" ]; then
    mkdir -p "${rootmnt}/writable/system-data/etc/systemd/system/systemd-hostnamed.service.d/"
    cat > "${rootmnt}/writable/system-data/etc/systemd/system/systemd-hostnamed.service.d/override.conf" << EOF
[Service]
PrivateNetwork=no
EOF
fi

mount -o bind "${rootmnt}/writable/system-data/etc/systemd/system" "${rootmnt}/etc/systemd/system"

## create "machine-id" if it does not exist and bind mount it (LP: #1619721)
if [ ! -e ${rootmnt}/writable/system-data/etc/machine-id ]; then
    cat /proc/sys/kernel/random/uuid|tr -d - >${rootmnt}/writable/system-data/etc/machine-id
fi
mount -o bind "${rootmnt}/writable/system-data/etc/machine-id" "${rootmnt}/etc/machine-id"

## create system-key
mkdir -p ${rootmnt}/writable/system-data/var/lib/snapd/
cat > ${rootmnt}/writable/system-data/var/lib/snapd/system-key << EOF
{"version":1}
EOF

## Pretend to have grub
mkdir -p "${rootmnt}/writable/system-data/boot/grub"
mount -o bind "${rootmnt}/writable/system-data/boot/grub" "${rootmnt}/boot/grub"
touch "${rootmnt}/writable/system-data/boot/grub/grub.cfg"
if [ ! -e  "${rootmnt}/writable/system-data/boot/grub/grubenv" ]; then
    chroot ${rootmnt} /usr/bin/grub-editenv /boot/grub/grubenv set snap_core=${SNAP_CORE} snap_kernel=${SNAP_KERNEL} snap_try_core= snap_try_kernel=
elif [ -n "${new_snap_mode}" ]; then
    chroot ${rootmnt} /usr/bin/grub-editenv /boot/grub/grubenv set snap_mode=${new_snap_mode}
fi

## Apply customized content
for user in "${rootmnt}"/writable/user-data/*
do
    if [ -d "${rootmnt}/custom/home" ] && [ ! -e "$user/.customized" ]; then
        echo "initrd: copying custom content tp " >/dev/kmsg || true
        cp -Rap "${rootmnt}"/custom/home/* "$user/"
        cp -Rap "${rootmnt}"/custom/home/.[a-zA-Z0-9]* "$user/"
        touch "$user/.customized"
        dstown=$(stat -c "%u:%g" "$user")
        chown -R "$dstown" "$user/"
    fi
done

## Workaround for nested LXD
mkdir -p "${rootmnt}/writable/system-data/var/lib/extrausers/"
echo "lxd:x:500:" >> "${rootmnt}/writable/system-data/var/lib/extrausers/group"

# Pivot to the new path
cd /root
pivot_root . mnt

# Move all the mounts
/bin/mount --move /mnt/dev /dev || true
/bin/mount --move /mnt/proc /proc || true
/bin/mount --move /mnt/sys /sys || true
/bin/umount -l /mnt

# Attempt to fix mount issues early
/bin/mount -a

# Exec new init
exec /sbin/init $@
