#!/bin/bash

# Variables declaration
declare -r BASE_PATH="/usr/share/volatile"
declare -r ARCHIVES_PATH="/usr/share/volatile/archives"
declare -r CASPER_PATH="/cdrom/casper"
declare -r RECOVERY_LABEL="Recovery Partition"
declare -a STAGES=("casper" "earlytask" "task" "latetask")

# Usage:
#  This function will touch a file as a flag for ${2} in this stage.
# params:
#  $1: stage; can be either 'casper', 'earlytask', 'task', 'latetask'
#  $2: action; can be any string
# returns:
#  0: success
#  254: failed to set
#  255: unknow argument
trigger_flag()
{
    local stage="$1"
    local action="$2"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    elif [ -z "$action" ]; then
        echo "The second argument should not be empty."
        return 255
    fi

    touch "${BASE_PATH}/${stage}-${action}" &&  return 0

    return 254
}

# Usage:
#  This function will check whether $2 flag is set or not in this stage.
# params:
#  $1: stage, can be either 'casper', 'earlytask', 'task', 'latetask'
#  $2: action; can be any string
# returns:
#  0: flag is unset
#  1: flag is set
#  255: unknow argument
check_flag()
{
    local stage="$1"
    local action="$2"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    elif [ -z "$action" ]; then
        echo "The second argument should not be empty."
        return 255
    fi

    test -f "${BASE_PATH}/${stage}-${action}" && return 1

    return 0
}

# Usage:
#  This function will clear reboot flag in this stage.
# params:
#  $1: stage, can be either 'casper', 'earlytask', 'task', 'latetask'
#  $2: action; can be any string
# returns:
#  0: success
#  254: failed to clean
#  255: unknow argument
clean_flag()
{
    local stage="$1"
    local action="$2"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    elif [ -z "$action" ]; then
        echo "The second argument should not be empty."
        return 255
    fi

    rm -f "${BASE_PATH}/${stage}-${action}" && return 0

    return 254
}

# Usage:
#  This function is for run some init commands in chroot stage
# returns:
#  0: success
chroot_task_init()
{
    # Move this line to a function if extra commands need to be executed
    create_package_index "${ARCHIVES_PATH}-task" || echo \
        "returns ($?); not found directory (255) or unable to cd to dir (254)"
    return 0
}

# Usage:
#  This function is for run some complete commands in chroot stage
# returns:
#  0: success
chroot_task_release()
{
    # Move this line to a function if extra commands need to be executed
    # TODO: 90_update_initramfs, remove to try
    ##remove the old dkms initrd to prevent the "Failed to find casper uuid.conf" error.
    #rm -f /boot/initrd.img-*.old-dkms
    # ---
    # TODO: 90_update_gru
    update-grub
    # TODO: 90_update_initramf, how about do nothing?
    depend -a
    update-initramfs -u
    return 0
}

# Usage:
#  This function is for run some init commands in binary stage
# returns:
#  0: success
binary_task_init()
{
    # Move this line to a function if extra commands need to be executed
    ubuntu-recovery-supports && mkdir -p "binary/debs/main"
    return 0
}

# Usage:
#  This function is for run some complete commands in binary stage
# returns:
#  0: success
binary_task_release()
{
    # Move this line to a function if extra commands need to be executed
    update_manifest_in_binary_stage
    return 0
}

# Usage:
#  This function is for run some init commands when ubiquity succeeded
#  Assuming be called by a script in scripts/ubiquity.d/success/
# returns:
#  0: success
ubiquity_succeeded_init()
{
    # Move this line to a function if extra commands need to be executed
    return 0
}

# Usage:
#  This function is for run some complete commands when ubiquity succeeded
#  Assuming be called by a script in scripts/ubiquity.d/success/
# returns:
#  0: success
ubiquity_succeeded_release()
{
    # Move this line to a function if extra commands need to be executed
    cp /var/log/volatile-task-casper.log /target/var/log
    return 0
}

# Usage:
#  This function is for run some init commands when ubiquity failed
#  Assuming be called by a script in scripts/ubiquity.d/failure/
# returns:
#  0: success
ubiquity_failed_init()
{
    # Move this line to a function if extra commands need to be executed
    return 0
}

# Usage:
#  This function is for run some complete commands when ubiquity failed
#  Assuming be called by a script in scripts/ubiquity.d/failure/
# returns:
#  0: success
ubiquity_failed_release()
{
    # Move this line to a function if extra commands need to be executed
    return 0
}

volatile_task_casper_init()
{
    # don't return event if any function got failed.
    return 0
}

volatile_task_casper_release()
{
    # don't return event if any function got failed.
    task_completed "casper" || echo \
        "failed to set task completed flag for earlytask"
    return 0
}

volatile_task_early_init()
{
    # don't return event if any function got failed.
    backup_sources_lists "/etc/apt" "earlytask" || echo \
        "failed to backup source list for earlytask"
    generate_source_list "earlytask" || echo \
        "returns $?; (255) unknow arguments; (254) no need to generate"
    set_oem_config_when_fed_enabled || echo \
        "failed to set oem config for fde for earlytask"
    return 0
}

volatile_task_early_release()
{
    # don't return event if any function got failed.
    restore_sources_lists "/etc/apt" "earlytask" || echo \
        "failed to restore source list for earlytask"
    remove_archives "earlytask" || echo \
        "failed to remove archives for earlytask"
    task_completed "earlytask" || echo \
        "failed to set task completed flag for earlytask"
    return 0
}

volatile_task_init()
{
    # don't return event if any function got failed.
    plymouth message \
        --text="Essential first boot setup is being performed. Please wait..."\
        || true
    backup_sources_lists "/etc/apt" "task" || echo \
        "failed to backup source list for task"
    mount_cdrom || echo \
        "returns $?; no need to mount cdrom (1) or mount failed (255)"
    generate_source_list "task" || echo \
        "returns $?; (255) unknow arguments; (254) no need to generate"
    set_oem_preseed
    return 0
}

volatile_task_release()
{
    # don't return event if any function got failed.
    umount_cdrom
    reset_machine_id
    # TODO: try without below message set
    #plymouth message --text=""
    restore_sources_lists "/etc/apt" "task" || echo \
        "failed to restore source list for task"
    remove_archives "task" || echo \
        "failed to remove archives for task"
    rename_recovery_bootefi || echo \
        "returns $?; no need to rename (1) or rename failed (255)"
    task_completed "task" || echo \
        "failed to set task completed flag for task"
    return 0
}

volatile_task_late_init()
{
    # FIXME: keep? make it more modulble
    pidof X >/dev/null && return 255 # origin is exit 255
    /sbin/plymouthd --mode=boot --attach-to-session || true
    plymouth show-splash || true
    plymouth message --text="Installation completed, rebooting ..." || true
    # end of 00_message

    return 0
}

volatile_task_late_release()
{
    # FIXME: make it as a function
    apt-get autoremove --purge --yes || true
    plymouth quit || true
    # end of 99_clean

    task_completed "latetask" || return $?
    return 0
}

# Usage:
#  This function is for run some general scripts in early oobe.
#  Assuming be called by a script in ubuntu/scripts/oem-config.d/early/
# returns:
#  0: success
oobe_init()
{
    # don't return event if any function got failed.
    mount_cdrom || echo \
        "returns $?; no need to mount cdrom (1) or mount failed (255)"
    return 0
}

# Usage:
#  This function is for run some general scripts in late oobe.
#  Assuming be called by a script in ubuntu/scripts/oem-config.d/late/
# returns:
#  0: success
oobe_release()
{
    # don't return event if any function got failed.
    umount_cdrom
    return 0
}

# Usage:
#  This function will run-part of specific directory, assuming be called by
#   a script under systemd service.
# params:
#  $1: stage, can be either 'casper', 'earlytask', 'task', 'latetask'
# returns:
#  0: success
#  1: not need to be executed
#  255: unknow argument
volatile_task_start()
{
    local stage="$1"
    local arg=""
    local scripts_dir=""
    local ret=0

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    fi

    # return if this stage already be ran.
    [ -e "${BASE_PATH}/${stage}-finished" ] && return 0

    # return if current stage is in installer
    [ "$stage" == "casper" ] && [ ! -d "$CASPER_PATH" ] && return 1
    [ "$stage" == "earlytask" ] && [ -d "$CASPER_PATH" ] && return 1
    [ "$stage" == "task" ] && [ -d "$CASPER_PATH" ] && return 1
    [ "$stage" == "latetask" ] && [ -d "$CASPER_PATH" ] && return 1

    exec >> /var/log/"volatile-${stage}".log 2>&1

    echo -n "volatile-${stage} starts at $(awk '{print $1}' /proc/uptime)"

    date

    # skip job if necessary
    if grep -q skip-"$stage" /proc/cmdline; then
        echo "volatile-${stage} job skipped"
        task_completed "$stage"
        return 1
    fi

    # make a fifo for debugging and wait for input to /tmp/pause
    if grep -q debug-"$stage" /proc/cmdline; then
        mkfifo -m 644 /tmp/pause; cat /tmp/pause; rm /tmp/pause;
    fi

    if [ "$stage" == "casper" ]; then
        arg="ubiquity"
        scripts_dir="${BASE_PATH}/ubiquity-starttask"
    elif [ "$stage" == "earlytask" ]; then
        arg="early"
        scripts_dir="${BASE_PATH}/${stage}"
    elif [ "$stage" == "task" ]; then
        arg="oem-config"
        scripts_dir="${BASE_PATH}/${stage}"
    elif [ "$stage" == "latetask" ]; then
        arg="late"
        scripts_dir="${BASE_PATH}/${stage}"
    else
        echo "'${stage}' doesn't in list."
        return 1
    fi

    [ "$stage" == "casper" ] && volatile_task_casper_init
    [ "$stage" == "earlytask" ] && volatile_task_early_init
    [ "$stage" == "task" ] && volatile_task_init
    [ "$stage" == "latetask" ] && volatile_task_late_init

    if [ -d "$scripts_dir" ]; then
        run-parts --report --arg="$arg" "$scripts_dir" 2>&1
    else
        echo "Not found customized scripts need to execute."
    fi

    [ "$stage" == "casper" ] && volatile_task_casper_release
    [ "$stage" == "earlytask" ] && volatile_task_early_release
    [ "$stage" == "task" ] && volatile_task_release
    [ "$stage" == "latetask" ] && volatile_task_late_release

    ret=0 && check_flag "$stage" "update-grub" || ret=$?
    if [ "$ret" -eq 1 ]; then
        clean_flag "$stage" "update-grub" || \
            echo "Failed to clean update grub flag"
        update-grub && echo "Failed to update grub"
    fi

    ret=0 && check_flag "$stage" "update-initramfs" || ret=$?
    if [ "$ret" -eq 1 ]; then
        clean_flag "$stage" "update-initramfs" || \
            echo "Failed to clean update initramfs flag"
        update-initramfs -u && echo "Failed to update initramfs"
    fi

    # TODO: here is a workaround for fixing LP: #1896210
    #  it needs to find out which stage sets recordfail and which stage
    #  that expects to unset it.
    [ "$stage" == "task" ] && grub-editenv /boot/grub/grubenv unset recordfail

    ret=0 && check_flag "$stage" "reboot" || ret=$?
    if [ "$ret" -eq 1 ]; then
        clean_flag "$stage" "reboot" || echo "Failed to clean reboot flag"
        plymouth message --text="Rebooting. Please wait..." || true
        echo "It's now rebooting for ${stage}..."

        ## TODO: we need to review it, below parts is came from sutton project
    	## wait until recordfail unset by /etc/init/grub-common,
    	## otherwise grub menu is shown and system could not be booted
    	## automatically (LP: #1219742)
    	#while grub-editenv /boot/grub/grubenv list | grep recordfail; do
    	#	sleep 1
    	#done

        # TODO: review below is what?
        /usr/share/ubuntu/scripts/anti-failure.sh reboot
    fi

    return 0
}

backup_sources_lists()
{
    local dir="${1:-/etc/apt}"
    local stage="$2"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-Second argument} should be either '${STAGES[*]}'"
        return 255
    fi

    for f in $(find "$dir" -type f -name "*.list" | grep "sources\.list" )
    do
        mv "$f" "${f%%.list}.volatile-${stage}"
    done

    return 0
}

generate_source_list()
{
    local stage="$1"
    local dir="${ARCHIVES_PATH}-${stage}"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'".
        return 255
    fi

    if [ -z "$dir" ] || [ ! -d "$dir" ]; then
        echo "Not found '${dir}' or '${dir}' is not a directory."
        return 254
    fi

    find "$dir"

    echo "deb file://${dir} /" > /etc/apt/sources.list

    # TODO: Needs review from stella (task/):
    # Supporting archives from Ubuntu Recovery partition. (LP: #1045684)
    if ubuntu-recovery-supports && mount | grep -q "cdrom"; then
        find "/cdrom"
        echo "deb file:///cdrom $(lsb_release -cs) main" >> \
            /etc/apt/sources.list
    fi

    # TODO: Needs review from sutton (task/):
    # XXX: a workaround?? still need?
    # in case Packages.gz is missing
    if [ -d "$dir" ] && [ ! -e "${dir}/Packages.gz" ]; then
        cd "$dir" && \
            apt-ftparchive packages . | gzip -n9 > Packages.gz
        cd - || true
    fi

    apt update

    return 0
}

restore_sources_lists()
{
    local dir="${1:-/etc/apt}"
    local stage="$2"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-Second argument} should be either '${STAGES[*]}'"
        return 255
    fi

    while read -r f; do
        mv "$f" "${f%%.volatile-${stage}}".list
    done < <(find "$dir" -type f -name "*.volatile-${stage}")

    rm -rf "/var/lib/apt/lists/*"

    return 0
}

remove_archives()
{
    local stage="$1"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    elif [ ! -e "${ARCHIVES_PATH}-${stage}" ]; then
        echo "${ARCHIVES_PATH}-${stage} not found."
        return 0
    fi

    rm -rf "${ARCHIVES_PATH}-${stage}"

    return 0
}

task_completed()
{
    local stage="$1"

    if [[ ! " ${STAGES[*]} " =~ (^|[[:space:]])$stage($|[[:space:]]) ]]; then
        echo "${stage:-First argument} should be either '${STAGES[*]}'"
        return 255
    fi

    touch "${BASE_PATH}/${stage}-finished"

    return 0
}

# Usage:
#  To check ubuntu recovery is supports or not.
# returns:
#  0: support.
#  1: not support.
ubuntu-recovery-supports()
{
    if blkid | grep -q "$RECOVERY_LABEL"; then
        return 0
    fi
    return 1
}

mount_cdrom()
{
    if ! ubuntu-recovery-supports; then
        return 1
    fi

    if ! mount | grep -q cdrom; then
        boot_dev="$(grub-probe -t device /)"
        [ -z "$boot_dev" ] && return 255
        root_dev="${boot_dev::-1}"
        if [ -d "/sys/firmware/efi" ]; then
            recovery_dev="${root_dev}2"
        else
            recovery_dev="${root_dev}1"
        fi
        mount -r "${recovery_dev}" /cdrom || return 255
    fi

    return 0
}

umount_cdrom()
{
    if mount | grep -q /cdrom; then
        umount /cdrom
    fi

    return 0
}

reset_machine_id()
{
    rm -f "/etc/machine-id"
    systemd-machine-id-setup
}

rename_recovery_bootefi()
{
    if ! ubuntu-recovery-supports; then
        return 1
    fi

    local RP MP
    RP=$(blkid | grep "$RECOVERY_LABEL" | awk -F: '{print $1}')
    MP=$(mktemp -d)

    if [ -z "$RP" ] || [ ! -d "$MP" ]; then
        echo "Not found partition '$RP' or directory '$MP'"
        return 255
    fi

    mount "$RP" "$MP" && \
        mv "$MP/efi/boot/bootx64.efi" "$MP/efi/boot/bootx64.efi.orig"
    umount "$MP"
    rm -rf "$MP"

    return 0
}

set_oem_preseed()
{
    if [ -f /cdrom/preseed/oem-preseed.cfg ]; then
        debconf-set-selections /cdrom/preseed/oem-preseed.cfg
    fi
    
    if [ -f /cdrom/preseed/ubuntu-recovery.cfg ]; then
        debconf-set-selections /cdrom/preseed/ubuntu-recovery.cfg
    fi

    return 0
}

create_package_index()
{
    local dir=$1

    if [ -z "$dir" ] || [ ! -d "$dir" ]; then
        echo "Not found '${dir}' or '${dir}' is not a directory."
        return 255
    fi

    cd "$dir" || return 254

    dpkg-scanpackages --multiversion . /dev/null > Packages
    cat Packages

    cd - || echo "Failed to return back to previous directory."

    return 0
}

# Usage:
#  Install all necessary packages provided by CE
# returns:
#  0: success
#  255: An error during installation
install_oem_packages()
{
    local opts="apt-get --yes --allow-downgrades \
--allow-remove-essential --allow-unauthenticated \
-o apt::install-recommends=true install"

    for pkg in $(ubuntu-drivers list| sort| awk '{print $1}' |sed 's/,//Ig'); do
        if [[ $pkg == oem-*-meta ]]; then
            echo "Found oem meta package: $pkg, installing.."
            eval "$opts $pkg" || continue
            factory_pkg=$(echo $pkg| sed 's/-/-factory-/2')
            echo "Try to install corresponding factory meta: $factory_pkg"
            eval "$opts --install-suggests $factory_pkg" || true
        fi
    done

    # Install other matched alias packages (e.g. official 8821ce, nvidia)
    ubuntu-drivers install

    return 0
}

# TODO: does it must be executed in volatile task?
update_manifest_in_binary_stage()
{
    #Add volatile-task packages to binary.mainfest
    chroot chroot/ dpkg-query -W \
        --showformat="\${Package}\\t\${Version}\\n" > \
        binary/casper/filesystem.manifest.chroot
    
    touch binary/casper/filesystem.manifest.volatile.and.langpacks
    
    while read -r x; do
        dpkg-deb -W "$x" >> \
            binary/casper/filesystem.manifest.volatile.and.langpacks
    done < \
        "$(find chroot/usr/share/volatile/ binary/pool -name "*.deb" -print)"
        
    cat binary/casper/filesystem.manifest.chroot \
        binary/casper/filesystem.manifest.volatile.and.langpacks \
        binary/casper/filesystem.manifest | sort -n | uniq > binary.manifest

    return 0
}

# FIXME: need to be review, early is necessary? late is necessary?
# it's a workaround or final solution?
set_oem_config_when_fed_enabled()
{
    # shellcheck source=/dev/null
    . /usr/share/debconf/confmodule

    local CRYPT CMD
    CRYPT=$(blkid | grep "crypto_LUKS")
    CMD=/usr/share/ubuntu/scripts/ubuntu-recovery-command

    # If there is a crypto_LUKS type of partition which means now is boot from
    # a bootable USB (to install system with FED feature, then set the oem-config
    # hooks manually (LP:1839427).
    if [ -n "$CRYPT" ]; then
        db_set oem-config/early_command $CMD oem-config early
        db_set oem-config/late_command $CMD oem-config late
    fi

    return 0
}
