#!/bin/bash
set -e -x

SCRIPTS_DIR=$(dirname ${0})

VERSION=${VERSION:?"VERSION not set"}

while getopts "i:f:c:d:t:r:o:p:ka:" OPTION
do
    case ${OPTION} in
        i) DIST="$OPTARG" ;;
        f) FILES="$OPTARG" ;;
        c) CLOUD_CONFIG="$OPTARG" ;;
        d) DEVICE="$OPTARG" ;;
        o) OEM="$OPTARG" ;;
        p) PARTITION="$OPTARG" ;;
        r) ROLLBACK_VERSION="$OPTARG" ;;
        k) KEXEC=y ;;
        a) APPEND="$OPTARG" ;;
        t) ENV="$OPTARG" ;;
        *) exit 1 ;;
    esac
done
[[ "$ARCH" == "arm" && "$ENV" != "rancher-upgrade" ]] && ENV=arm

DIST=${DIST:-/dist}
CLOUD_CONFIG=${CLOUD_CONFIG:-"${SCRIPTS_DIR}/conf/empty.yml"}
CONSOLE=tty0
BASE_DIR="/mnt/new_img"
BOOT=boot/
# TODO: Change this to a number so that users can specify.
# Will need to make it so that our builds and packer APIs remain consistent.
PARTITION=${PARTITION:=${DEVICE}1}
KERNEL_ARGS="rancher.state.dev=LABEL=RANCHER_STATE rancher.state.wait console=${CONSOLE}"

device_defined()
{
    if [[ -z "$1" ]]; then
        echo "Need to Pass a device name -d <dev>." 1>&2
        exit 1
    fi
}

format_device()
{
    device_defined ${DEVICE}
    mkfs.ext4 -F -i 4096 -L RANCHER_STATE ${PARTITION}
}

get_dev()
{
    if [ -z "$(which ros)" ]; then
        lsblk -n -o label $1
    else
        ros dev LABEL=${1}
    fi
}

mount_device()
{
    LABEL=RANCHER_STATE
    local raw="${1:-false}"

    mkdir -p ${BASE_DIR}

    if [ -n "$(get_dev RANCHER_BOOT)" ]; then
        LABEL=RANCHER_BOOT
        BOOT=
    fi

    local mount_opts="-L ${LABEL}"
    if [ "${raw}" == "true" ]; then
        device_defined ${DEVICE}
        mount_opts=${PARTITION}
    fi

    mount ${mount_opts} ${BASE_DIR}
    trap "umount ${BASE_DIR}" EXIT
}

create_boot_dirs()
{
    mkdir -p ${BASE_DIR}/${BOOT}grub
    mkdir -p ${BASE_DIR}/${BOOT}syslinux
}

install_syslinux() {
    dd bs=440 count=1 if=/usr/lib/syslinux/mbr/mbr.bin of=${DEVICE}
    cp /usr/lib/syslinux/modules/bios/* ${BASE_DIR}/${BOOT}syslinux
    extlinux --install ${BASE_DIR}/${BOOT}syslinux
}

install_syslinux_raid() {
    dd bs=440 count=1 if=/usr/lib/syslinux/mbr/mbr.bin of=/dev/sda
    dd bs=440 count=1 if=/usr/lib/syslinux/mbr/mbr.bin of=/dev/sdb
    cp /usr/lib/syslinux/modules/bios/* ${BASE_DIR}/${BOOT}syslinux
    extlinux --install --raid ${BASE_DIR}/${BOOT}syslinux
}

install_grub() {
    grub-install --boot-directory=${BASE_DIR}/boot ${DEVICE}
}

grub2_config(){
    local grub_cfg=${BASE_DIR}/${BOOT}grub/grub.cfg
    local append_line="${1}"
cat >${grub_cfg} <<EOF
set default="0"
set timeout="1"
#set fallback=1

menuentry "RancherOS-current" {
  set root=(hd0,msdos1)
  linux /${BOOT}vmlinuz-${VERSION}-rancheros ${KERNEL_ARGS} ${append_line}
  initrd /${BOOT}initrd-${VERSION}-rancheros
}

EOF

if [ ! -z ${ROLLBACK_VERSION} ]; then
    sed -i 's/^#set/set/' ${grub_cfg}
cat >>${grub_cfg} <<EOF
menuentry "RancherOS-rollback" {
  set root=(hd0,msdos1)
  linux /${BOOT}vmlinuz-${ROLLBACK_VERSION}-rancheros ${KERNEL_ARGS} ${append_line}
  initrd /${BOOT}initrd-${ROLLBACK_VERSION}-rancheros
}
EOF
fi

}

syslinux_config(){
    local syslinux_cfg=${BASE_DIR}/${BOOT}syslinux/syslinux.cfg
    local append_line="${1}"
cat >${syslinux_cfg} <<EOF
DEFAULT RancherOS-current

LABEL RancherOS-current
    LINUX ../vmlinuz-${VERSION}-rancheros
    APPEND ${KERNEL_ARGS} ${append_line}
    INITRD ../initrd-${VERSION}-rancheros

EOF

if [ ! -z ${ROLLBACK_VERSION} ]; then
cat >>${syslinux_cfg} <<EOF
LABEL RancherOS-rollback
    LINUX ../vmlinuz-${ROLLBACK_VERSION}-rancheros
    APPEND ${KERNEL_ARGS} ${append_line}
    INITRD ../initrd-${ROLLBACK_VERSION}-rancheros

EOF
fi

}

install_rancher()
{
    cp ${DIST}/initrd ${BASE_DIR}/${BOOT}initrd-${VERSION}-rancheros
    cp ${DIST}/vmlinuz ${BASE_DIR}/${BOOT}vmlinuz-${VERSION}-rancheros
}

pvgrub_config()
{
    local grub_file=${BASE_DIR}/${BOOT}grub/menu.lst
    local append_line="${1}"
cat > ${grub_file}<<EOF
default 0
timeout 0
#fallback 1

hiddenmenu

title RancherOS ${VERSION}-(current)
root (hd0)
kernel /${BOOT}vmlinuz-${VERSION}-rancheros ${KERNEL_ARGS} ${append_line}
initrd /${BOOT}initrd-${VERSION}-rancheros

EOF

if [ ! -z ${ROLLBACK_VERSION} ]; then
    sed -i 's/^#\(fallback\)/\1/' ${grub_file}
cat >> ${grub_file}<<EOF
title RancherOS ${ROLLBACK_VERSION}-(rollback)
root (hd0)
kernel /${BOOT}/vmlinuz-${ROLLBACK_VERSION}-rancheros ${KERNEL_ARGS} ${append_line}
initrd /${BOOT}initrd-${ROLLBACK_VERSION}-rancheros
EOF
fi
}

format_and_mount()
{
    format_device
    mount_device
    create_boot_dirs
}

if [ -n ${ENV} ]; then
    case ${ENV} in
        "generic")
            format_and_mount
            install_grub
            "${SCRIPTS_DIR}/seed-data" ${BASE_DIR} ${CLOUD_CONFIG} ${FILES}
            ;;
        "arm")
            format_and_mount
            "${SCRIPTS_DIR}/seed-data" ${BASE_DIR} ${CLOUD_CONFIG} ${FILES}
            ;;
        "amazon-ebs-pv"|"amazon-ebs-hvm")
            CONSOLE=ttyS0
            format_and_mount
            if [ "${ENV}" == "amazon-ebs-hvm" ]; then
                install_grub
            fi
            # AWS Networking recommends disabling.
            "${SCRIPTS_DIR}/seed-data" ${BASE_DIR} ${CLOUD_CONFIG} ${FILES}
            ;;
        "googlecompute")
            CONSOLE=ttyS0
            format_and_mount
            install_grub
            "${SCRIPTS_DIR}/seed-data" ${BASE_DIR} ${CLOUD_CONFIG} ${FILES}
            ;;
        "noformat")
            mount_device
            create_boot_dirs
            install_syslinux
            ;;
        "raid")
            mount_device
            create_boot_dirs
            install_syslinux_raid
            ;;
        "bootstrap")
            CONSOLE=ttyS0
            mount_device true
            create_boot_dirs
            KERNEL_ARGS="${KERNEL_ARGS} rancher.cloud_init.datasources=[ec2,gce]"
            ;;
        "rancher-upgrade")
            mount_device
            create_boot_dirs
            ;;
        *)
            echo "$ENV is not a valid environment" 1>&2
            exit 1
            ;;
    esac
fi

if [ -e ${BASE_DIR}/${BOOT}append ]; then
    PRESERVED_APPEND=$(cat ${BASE_DIR}/${BOOT}append)
fi
if [ "${APPEND}" = "" ]; then
    APPEND="${PRESERVED_APPEND}"
fi
echo "${APPEND}" > ${BASE_DIR}/${BOOT}append

grub2_config "${APPEND}"
syslinux_config "${APPEND}"
pvgrub_config "${APPEND}"
install_rancher

seusers=${BASE_DIR}/etc/selinux/ros/seusers
failsafe_context=${BASE_DIR}/etc/selinux/ros/contexts/failsafe_context
if [ -f "${seusers}" ]; then
    echo "__default__:unconfined_u:s0-s0:c0.c1023" > ${seusers}
fi
if [ -f "${failsafe_context}" ]; then
    echo "unconfined_r:unconfined_t:s0" > ${failsafe_context}
fi

if [ "$KEXEC" = "y" ]; then
    kexec -l ${DIST}/vmlinuz --initrd=${DIST}/initrd --append="${KERNEL_ARGS} ${APPEND}" -f
fi