#!/bin/bash

TMP_LOCATION="/var/tmp/ubuntu-upgrade-testing"
BASE_LOCATION="${ADT_ARTIFACTS}/upgrade_run_config"
PRE_SCRIPT_LOCATION="${BASE_LOCATION}/pre_scripts"
POST_SCRIPT_LOCATION="${BASE_LOCATION}/post_scripts"
# Currently experimenting with using yaml output for run results (test
# pass/fail etc.) for now, then we'll use something better
TEST_RESULTS_DIR="${ADT_ARTIFACTS}/upgrade_run"
TEST_RESULT_FILE="${TEST_RESULTS_DIR}/runner_results.yaml"
CANARY_NAME="/tmp/upgrade_script_reboot_canary"

# Only copy on the first run through
if [ ! -d "${BASE_LOCATION}" ]; then
   mkdir "${BASE_LOCATION}"
   mv "${TMP_LOCATION}/pre_scripts" "${BASE_LOCATION}"
   mv "${TMP_LOCATION}/post_scripts" "${BASE_LOCATION}"
   mv "${TMP_LOCATION}/auto_upgrade_test_settings" "${BASE_LOCATION}"
fi

# This is put in a known place by the wrapper script and contains the details
# of above (as they will change each run).
CONFIG_FILE="${BASE_LOCATION}/auto_upgrade_test_settings"
source "${CONFIG_FILE}"
export TEST_RESULTS_DIR

HAVE_REBOOTED=$ADT_REBOOT_MARK

STATUS=0

function upgrade_log() {
    local output=$1
    echo -e "auto-upgrade [$(date +%R:%S)]: ${output}"
}

function cleanup() {
    # Collect the results at exit so we cover both successful runs and
    # failures.
    collect_results
    cp "${CONFIG_FILE}" "${TEST_RESULTS_DIR}"
    upgrade_log "Cleaning up configuration files."
    rm -r "${BASE_LOCATION}"
}

function main() {
    # Ensure we don't have any mix-ups with multiple runs on the same testbed.
    trap cleanup EXIT

    upgrade_log "Running on ${RUNNING_BACKEND}"

    if [ -z "${HAVE_REBOOTED}" ]; then
        upgrade_log "Beginning from the start."
        create_reboot_canary

        output_running_system

        do_setup

        exit_if_not_running_initial_system

        pre_tests
        STATUS=$?
        exit_with_log_if_nonzero $STATUS "ERROR: Something went during the prerun scripts."

        store_prereboot_details
        do_upgrade_and_maybe_reboot
    else
        upgrade_log "Skipping pre-tests as we have rebooted."
    fi

    # If we have rebooted we pick up from here.
    output_running_system
    exit_if_reboot_canary_exists
    exit_if_havent_upgraded

    # Check if we need to do another upgrade/reboot
    if need_another_upgrade; then
        echo "Appears we're in a multi-part upgrade. Upgrading/rebooting again."
        create_reboot_canary
        do_upgrade_and_maybe_reboot
    else
        exit_if_not_running_expected_post_system

        # No need to explicitly exit here as we're at the end.
        post_tests
        STATUS=$?
    fi

    exit $STATUS
}

function exit_with_log_if_nonzero() {
    local retcode=$1
    local error_message=$2
    if (( retcode != 0 )); then
        upgrade_log "ERROR: ${error_message}"
        exit 1
    fi
}

function exit_if_not_running_initial_system() {
    local running_system=$(_get_running_system_name)
    upgrade_log "Checking that running system (${running_system}) is ${INITIAL_SYSTEM_STATE}"
    if [ "${INITIAL_SYSTEM_STATE}" != "${running_system}" ]; then
        upgrade_log "ERROR: Expected ${INITIAL_SYSTEM_STATE} got ${running_system}"
        # Is there a better way than just exiting here?
        exit 1
    fi
}

# Can we de-dupe these methods too?
function exit_if_not_running_expected_post_system() {
    local running_system=$(_get_running_system_name)
    upgrade_log "Checking that running system (${running_system}) is ${POST_SYSTEM_STATE}"
    if [ "${POST_SYSTEM_STATE}" != "${running_system}" ]; then
        upgrade_log "ERROR: Expected ${POST_SYSTEM_STATE} got ${running_system}"
        # Is there a better way than just exiting here?
        exit 1
    fi
}

function exit_if_havent_upgraded() {
    local running_system_version=$(get_current_version)
    upgrade_log "Checking that an upgrade has occured."
    if [ "${BEFORE_REBOOT_VERSION}" != "" ] && [ "${running_system_version}" == "${BEFORE_REBOOT_VERSION}" ]; then
        upgrade_log "ERROR: Still the same system version after reboot"
        exit 1
    fi
}

function create_reboot_canary() {
    touch "${CANARY_NAME}"
}

function store_prereboot_details() {
    # Store details that we'll use after a reboot.
    # Current running version as we way need to reboot between versions.
    echo "BEFORE_REBOOT_VERSION=$(get_current_version)" >> "${CONFIG_FILE}"
}

function exit_if_reboot_canary_exists() {
    if [ -f "${CANARY_NAME}" ]; then
        upgrade_log "ERROR: system has not rebooted"
        exit 1
    fi
}

function _get_running_system_name() {
    echo $(lsb_release -sc)
}

function pre_tests() {
    # Script setup and run. For each test:
    #  - create a output dir for the results and make available to script
    #  - Run script
    #  - Log success or failure of script
    echo "pre_script_output:" >> ${TEST_RESULT_FILE}
    success=0
    for test in $PRE_TESTS_TO_RUN; do

        local this_script_results="${TEST_RESULTS_DIR}/pre_${test}/"
        mkdir "${this_script_results}"
        export TESTRUN_RESULTS_DIR=$this_script_results

        local FULL_TEST_SCRIPT_PATH="${PRE_SCRIPT_LOCATION}/${test}"
        upgrade_log "Running test: ${FULL_TEST_SCRIPT_PATH} -- Results: ${this_script_results}"
        ${FULL_TEST_SCRIPT_PATH}

        local test_result=$?
        if (( test_result != 0 )); then
            echo "  \"${test}\": FAIL" >> ${TEST_RESULT_FILE}
            success=1
        else
            echo "  \"${test}\": PASS" >> ${TEST_RESULT_FILE}
        fi
    done
    return $success
}

function post_tests() {
    # Script setup and run. For each test:
    #  - create a output dir for the results and make available to script
    #  - Run script
    #  - Log success or failure of script
    echo "post_test_output:" >> $TEST_RESULT_FILE
    success=0
    for test in $POST_TESTS_TO_RUN; do
        local this_script_results="${TEST_RESULTS_DIR}/post_${test}/"
        mkdir "${this_script_results}"
        export TESTRUN_RESULTS_DIR=$this_script_results

        local FULL_TEST_SCRIPT_PATH="${POST_SCRIPT_LOCATION}/${test}"
        upgrade_log "Running test: ${FULL_TEST_SCRIPT_PATH} -- Results: ${this_script_results}"

        ${FULL_TEST_SCRIPT_PATH}

        local test_result=$?
        if (( test_result != 0 )); then
            echo "  \"${test}\": FAIL" >> $TEST_RESULT_FILE
            success=1
        else
            echo "  \"${test}\": PASS" >> $TEST_RESULT_FILE
        fi
    done
    return $success
}

function do_setup() {
    upgrade_log "Performing run setup."
    # Make sure the output results file is available and proper yaml.
    mkdir "${TEST_RESULTS_DIR}"
    echo "---" >> "${TEST_RESULT_FILE}"
}

function need_another_upgrade() {
    # Check if we're not running the right version
    # If not are we able to upgrade to the right version?
    local running_system=$(_get_running_system_name)
    if [ "${POST_SYSTEM_STATE}" != "${running_system}" ]; then
        potential_upgrade_version=$(get_potential_upgrade_version)
        current_version=$(get_current_version)
        echo "Comparing ${potential_upgrade_version} against ${current_version}"
        # we can upgrade further and the upgrade target is greater than our current system.
        if [ "${potential_upgrade_version}" ] && version_lt "${current_version}" "${potential_upgrade_version}"; then
            return 0
        fi
    fi
    return 1
}

function get_current_version() {
    echo $(lsb_release -rs)
}

function get_potential_upgrade_version() {
    # Attempt to get the version that we would upgrade to. Attempts to use
    # development version if needed.
    # Might return an empty string if there are no upgrade candidates at all.
    local version=$(do-release-upgrade -p -c | awk '/New release/ {print $3}' | tr -d \')
    if [ ! "${version}" ]; then
        # Lets try for a development version
        local version=$(do-release-upgrade -c -d | awk '/New release/ {print $3}' | tr -d \')
        echo "${version}"
    else
        echo "${version}"
    fi
}

# version_lte and version_lt taken from: http://stackoverflow.com/a/4024263
function version_lte() {
    [  "$1" = "`echo -e "$1\n$2" | sort --version-sort | head -n1`" ]
}

function version_lt() {
    [ "$1" = "$2" ] && return 1 || version_lte $1 $2
}

function do_upgrade_and_maybe_reboot() {
    current="${INITIAL_SYSTEM_STATE}"
    target="${POST_SYSTEM_STATE}"
    upgrade_log "Attempting to upgrade from ${current} to ${target}"

    do_normal_upgrade
    exit_with_log_if_nonzero $STATUS "ERROR: Something went wrong with the upgrade."
    maybe_reboot

    exit_with_log_if_nonzero $STATUS "ERROR: Something went wrong with the upgrade."

    upgrade_log "Upgrading complete."
}

function have_internet_connection() {
    ping_count=0
    max_ping=5
    while  ! ping -c 1 launchpad.net > /dev/null 2>&1 && ((ping_count < max_ping)); do
        update_log "Failed to ping launchpad.net. Seems there is no Internet connection"
        ((ping_count++))
        sleep 1
    done
    echo "Ping count: ${ping_count}"
    if ((ping_count==max_ping)); then
        return 1
    else
        return 0
    fi

}

function reboot_prepare() {
    prepare_function="/tmp/autopkgtest-reboot-prepare"
    if [ -f ${prepare_function} ]; then
        upgrade_log "Preparing the system for reboot."
        $($prepare_function 'upgradetests')
        sleep 30
    else
        upgrade_log "This testbed does not support rebooting."
        exit 1
    fi
}

function do_normal_upgrade() {
    upgrade_log "Starting machine upgrade."

    export DEBIAN_FRONTEND=noninteractive
    # Ensure we have do-release-upgrade
    apt-get update
    apt-get -y dist-upgrade
    apt-get -y install openssh-server update-manager-core

    # Allow upgrade from lts to non-lts if there's not lts to upgrade to
    local version
    if grep '^Prompt=lts' /etc/update-manager/release-upgrades; then
        version=$(do-release-upgrade -d -p -c | awk '/New release/ {print $3}' | tr -d \')
        if [ -z "${version}" ]; then
            # No LTS release to upgrade to. Enable non-LTS upgrades.
            sed -i 's/Prompt=lts/Prompt=normal/' /etc/update-manager/release-upgrades
        fi
    fi

    version=$(do-release-upgrade -p -c | awk '/New release/ {print $3}' | tr -d \')
    if [ -z "${version}" ]; then
        do-release-upgrade -d -f DistUpgradeViewNonInteractive
    else
        do-release-upgrade -p -f DistUpgradeViewNonInteractive
    fi

    STATUS=$?
}

function maybe_reboot() {
    # Check if we actually want to reboot . . .
    reboot_function="/tmp/autopkgtest-reboot"
    if [ -f ${reboot_function} ]; then
        upgrade_log "Rebooting the system."
        if [ "${RUNNING_BACKEND}" = "lxc" ]; then
            # lxc reboot is doing something different to expected.
            rm "${CANARY_NAME}"
        fi
        $($reboot_function 'upgradetests')
    else
        upgrade_log "This testbed does not support rebooting."
        exit 1
    fi
}

function collect_results() {
    # Move any files of interest into $TEST_RESULTS_DIR
    upgrade_log "Collecting system details."
    system_details_dir="${TEST_RESULTS_DIR}/system_details"
    mkdir "${system_details_dir}"
    cp -fr /var/log/dist-upgrade "${system_details_dir}/dist-upgrade/"
    cp /var/log/dpkg.log "${system_details_dir}/"
    cp -fr /etc/apt/ "${system_details_dir}/apt/"
}

function output_running_system() {
    echo "Currently running: $(lsb_release -a)"
}

main "$@"
