#!/usr/bin/env bash
# This script is run after the first reboot by zoox_first_boot_script which is
# created by zoox_init.sh. This script runs setup script for both ai and
# vision machines after reboot (reboot is needed to get python semaphores
# working which work on a multi user run level mode). This script is only runs
# on main-ai machine. This script takes 5 parameters i.e.,
# 1. Devuser i.e., zoox
# 2. Robot name
# 3. Total number machines in the l3 environment.
# 4. Driving branch to be used for setup script.
# 5. Compute version: Compute major version
# 6. Compute version: Compute minor version
# 7. Operator Version: Version of the operator
# 8. Robot has operator: 1 to indicate that operator exists and can be imaged as part of robot otherwise 0

# Note: Explicitly not setting a set -eou flag as this script should not return
# error
exec >> /var/log/zooxpxeinstall.log 2>&1

devuser="$1"
robot_name="$2"
vision_count=$(expr $3 - 1)
driving_sha="$4"
compute_major_version="$5"
compute_minor_version="$6"
operator_version="$7"
robot_has_operator="$8"

ki_robot_info_url="{{ robot_info_base_url }}/v1/robot/?name=${robot_name}"
# This variable is going to be populated by
robot_network_type=""

if [[ "${robot_has_operator}" == "1" ]];
then
    vision_count=$(expr ${vision_count} - 1)
fi

nc_timeout=5
nc_check_port=22
slack_url="{{ cobbler_slack_url }}"

# Max number of check to be performed to check if visions/operator machines are up and running.
max_number_of_checks=20
sleep_duration_between_vision_checks=30
slack_upload_url="{{ cobbler_slack_upload_url }}"
slack_upload_channel="{{ cobbler_slack_upload_channel }}"
slack_upload_token="{{ cobbler_slack_upload_token }}"
firefox_settings_url="{{ firefox_settings_url }}"
firefox_settings_path="/var/tmp/firefox.tar.gz"
firefox_expected_sha="{{ firefox_settings_sha }}"

ki_status_url="{{ cobbler_ki_status_url }}imaging_status/"

desktop_image_base_url="{{ robot_images_base_url }}"

kernel_version=$(uname -r)
imaging_server="{{ inventory_hostname }}"
mac_addresses=$(ip link | grep 'link/ether' | awk '{print $2}' | tr '\n' ',')

imageries_tar_name="imageries.tar.gz"
imageries_tar_path="/var/tmp/${imageries_tar_name}"

chown ${devuser}:${devuser} /var/log/zooxpxeinstall.log

function wait_for_network() {
    state="error";

    while [[ "${state}" == "error" ]]; do
        #do a ping and check that its not a default message or change to grep for something else
        state=$(ping -q -w 1 -c 1 {{ cobbler_dhcp_router_ip }} > /dev/null && echo ok || echo error)

        #sleep for 2 seconds and try again
        sleep 2
    done

}

function log_message() {
    # updates motd and sends a slack message
    # sends message to slack channel provided by slack_url
    message="$1"
    status="$2"

    status_text="SUCCESS"

    if [[ ${status} -ne 0 ]];
    then

        echo "#!/bin/sh
echo \"
############################################################
# ZOOX: DO NOT USE THIS MACHINE. IT IS NOT IMAGED PROPERLY #
############################################################

Error message is: $1
\"" >> /etc/update-motd.d/99-zoox-image-status

        status_text="FAILED"
    fi

    curl  --verbose --request POST \
    --header "Content-Type: application/json" \
    --header "Accept: application/json"  \
    --data "{\"robot\":\"${robot_name}\",\"status\":\"${status_text}\",\"imagingServer\":\"${imaging_server}\",\"linuxKernelVersion\":\"${kernel_version}\",\"MACAddresses\":\"${mac_addresses}\",\"drivingSHA\":\"${driving_sha}\"}" \
    ${ki_status_url}

    curl --verbose --request POST --data-urlencode "payload={\"text\":\"$message\",\"username\":\"$robot_name\"}" ${slack_url}

    curl --silent --output /dev/null --form token="$slack_upload_token" --form channels="$slack_upload_channel" --form filename="$robot_name-$(date +%Y-%m-%d:%H:%M:%S)-log.txt" --form file=@/var/log/zooxpxeinstall.log  ${slack_upload_url}
}

function update_robot_network_type() {
    robot_network_type=$(curl --silent "${ki_robot_info_url}" | jq ".items[0].robotNetworkType" | sed "s/\"//g")

    echo "The robot ${robot} has ${robot_network_type} as robot network type"
}

function setup_group() {
    host_group="$1"
    is_limited_group=1
    retry_count=0
    max_retries=2

    sim_tag=""
    unsafe_tag="--unsafe"
    limit_hosts=""
    skip_tags_for_limited=""
    imaging_ansible_envs="ZOOX_ANSIBLE_INVENTORY_ROBOT_NAME=${robot_name}"
    imaging_ansible_envs+=" ZOOX_ANSIBLE_SKIP_INACTIVE_VISION_HOSTS=True"
    imaging_ansible_envs+=" ZOOX_VISION_COUNT=${vision_count}"

    # Add sim flag if robot is zimbeast
    if [[ "${robot_name}" == *"zimbeast"* ]];
    then
        sim_tag="--sim"
    fi

    if [[ "${host_group}" == "vision_limited" ]];
    then
        # don't run unsafe tag for vision as it would mess up networking and main-ai would run it
        unsafe_tag=""
        # limit the host to only current vision host
        limit_hosts="--limit=$(hostname)"
        # reset the host group for vision limited
        host_group="vision"
        is_limited_group=0
        skip_tags_for_limited=",ssh_keys_vision"
    fi

    if [[ "${host_group}" == "operator_limited" ]];
    then
        # don't run unsafe tag for operator as it would mess up networking and main-ai would run it
        unsafe_tag=""
        # reset the host group for operator limited
        host_group="operator"
        is_limited_group=0
        skip_tags_for_limited=""
    fi

    while [[ ${retry_count} -lt ${max_retries} ]];
    do

        sudo su ${devuser} --command "cd /home/$devuser/driving && ${imaging_ansible_envs} \
        ./scripts/setup/setup.sh ${host_group} --no-ask-pass ${unsafe_tag} ${sim_tag} --nvidia \
        {% if cobbler_setup_skip_tags|length > 0 -%}--skip-tags='{{ cobbler_setup_skip_tags|join(',') }}{%- endif %}${skip_tags_for_limited}' \
        --extra-vars=\"driving_sha=${driving_sha}\" \
        --extra-vars=\"compute_major_version=${compute_major_version}\" \
        --extra-vars=\"compute_minor_version=${compute_minor_version}\" \
        --extra-vars=\"robot_network_type=${robot_network_type}\" \
        --extra-vars=\"operator_version=${operator_version}\" \
        --extra-vars=\"image_mode=yes\" \
        ${limit_hosts}"
        result="$?"

        if [[ ${result} -eq 0 ]];
        then
            break;
        fi

        retry_count=$(expr ${retry_count} + 1)
        echo "############ setup $host_group on $robot_name failed, retry count is  $retry_count ###########"

    done

    if [[ ${result} -ne 0 && ${is_limited_group} -ne 0 ]];
    then
        log_message "Failed! Unable to setup $host_group on $robot_name" ${result}
        exit ${result}
    fi
}

function ensure_vision_hosts_and_setup() {

    if [[ ${vision_count} -eq 0 ]];
    then
        return
    fi

    number_of_checks=0

    while true;
    do
        visible_visions_count=0
        missing_visions=""
        first_missing_vision="false"
        for i in $(eval echo "{1..${vision_count}}");
        do
            vision_machine_name="vision$i"
            nc -zw "${nc_timeout}" "${vision_machine_name}" ${nc_check_port}
            vision_exists="$?"

            if [[ ${vision_exists} -eq 0 ]]; then
                visible_visions_count=$(expr ${visible_visions_count} + 1)
            else
                if [[ "${first_missing_vision}" == "false" ]]; then
                    missing_visions="${vision_machine_name}"
                    first_missing_vision="true"
                else
                    missing_visions="${missing_visions}, ${vision_machine_name}"
                fi

            fi
        done

        if [[ ${visible_visions_count} -eq ${vision_count} ]];
        then
            break;
        fi

        number_of_checks=$(expr ${number_of_checks} + 1)

        if [[ ${number_of_checks} -ge ${max_number_of_checks} ]];
        then
            message="Failed! Expected number of vision machines are not up. Missing visions are '${missing_visions}'"
            log_message "${message}"  -1
            exit 1
        fi
        sleep ${sleep_duration_between_vision_checks};
    done

    # Visions are up but lets wait for visions to be ready.
    sleep 500

    setup_group "vision"

    # Reboot vision nodes
    for i in $(eval echo "{1..${vision_count}}");
    do
        sudo su ${devuser} --command "ssh vision$i \"sudo reboot\""
    done
}

function ensure_operator_up() {
    if [[ "${robot_has_operator}" == "0" ]];
    then
        return
    fi

    number_of_checks=0
    while [[ ${number_of_checks} -lt ${max_number_of_checks} ]];
    do
        nc -zw "$nc_timeout" "operator" ${nc_check_port}
        operator_is_up="$?"

        if [[ ${operator_is_up} -eq 0 ]];
        then
            break;
        fi

        number_of_checks=$(expr ${number_of_checks} + 1)

    done

    if [[ ${number_of_checks} -ge ${max_number_of_checks} ]];
    then
        log_message "Failed! Operator machine is not up" -1
        exit 1
    fi
}

function ensure_operator_up_and_setup() {
    ensure_operator_up

    if [[ ${vision_count} -eq 0 ]];
    then
        # Operator is up but wait for sometime to make sure that operator is ready.
        sleep 300
    fi

    setup_group "operator"
}

function reboot_operator() {
    if [[ "${robot_has_operator}" == "0" ]];
    then
        return
    fi

    # Reboot operator
    sudo su ${devuser} --command "ssh operator \"sudo reboot\""
}

function download_imageries() {
    sudo su ${devuser} --command "wget http://{{ cobbler_dhcp_router_ip }}/${imageries_tar_name} --quiet --output-document=${imageries_tar_path} || rm --force ${imageries_tar_path}"

    # Check if the tar file exist and extract it.
    if [[ -f "${imageries_tar_path}" ]];
    then
        echo "Extracting the imageries tar $(date)"
        sudo su ${devuser} --command "tar --extract --file ${imageries_tar_path} --directory=/home/${devuser}/"
        echo "Removing the imageries tar $(date)"
        rm --force "${imageries_tar_path}"
    fi

    if [[ "${robot_has_operator}" == "0" ]];
    then
        echo "Robot has no operator"
        return
    fi


    # copy imageries to operator (only the 2d map)
    number_of_checks=0

    file_list_file=$(mktemp)
    for d in /home/"${devuser}"/imageries/*
    do
        if [ -d "$d/images/laser" ]
        then
            echo "$(basename $d)/images/laser/" >> "$file_list_file"
        fi
    done
    sudo chmod a+r ${file_list_file}

    while [[ ${number_of_checks} -lt ${max_number_of_checks} ]];
    do
        echo "Syncing imageries to operator"

        sudo su ${devuser} --command "rsync --rsh='ssh' --files-from=${file_list_file} /home/${devuser}/imageries ${devuser}@operator:/home/${devuser}/imageries"
        operator_is_synced="$?"

        if [[ ${operator_is_synced} -eq 0 ]];
        then
            echo "Synced imageries to operator successfully"
            # For debugging, log the size of the imagery on the operator machine.
            sudo su ${devuser} --command "ssh ${devuser}@operator 'du -sh /home/${devuser}/imageries/*'"
            break;
        fi

        number_of_checks=$(expr ${number_of_checks} + 1)

    done


}

wait_for_network

sudo apt-get update

# Setup workspace switcher
sudo su ${devuser} --command "dbus-launch gsettings set org.compiz.core:/org/compiz/profiles/unity/plugins/core/ hsize 2"
sudo su ${devuser} --command "dbus-launch gsettings set org.compiz.core:/org/compiz/profiles/unity/plugins/core/ vsize 2"

# Setup infinite terminal scrolling
sudo su ${devuser} --command "gconftool-2 -t bool --set /apps/gnome-terminal/profiles/Default/scrollback_unlimited 1"

# copy firefox settings
wget "${firefox_settings_url}" --output-document=${firefox_settings_path}

actual_sha="$(sha256sum "${firefox_settings_path}" | awk '{print $1}')"

if [[ ${actual_sha} == ${firefox_expected_sha} ]];
then
    mkdir /home/${devuser}/.mozilla
    tar -xf /var/tmp/firefox.tar.gz --directory=/home/${devuser}/.mozilla/
    chown -R ${devuser}:${devuser} /home/${devuser}/.mozilla
else
    echo "The sha did not match."
fi

# Update the desktop background
desktop_image_url="${desktop_image_base_url}"

if [[ $(hostname) == "operator" ]];
then
    desktop_image_url="${desktop_image_url}/operator.jpg"

elif [[ $(hostname) =~ ^vision.* ]];
then
    desktop_image_url="${desktop_image_url}/$(hostname).jpg"
else
    desktop_image_url="${desktop_image_url}/main-ai.jpg"
fi

if [[ "${desktop_image_url}" != "${desktop_image_base_url}" ]];
then
    mkdir /home/${devuser}/Pictures
    wget "${desktop_image_url}" --output-document=/home/${devuser}/Pictures/desktop.jpg
    chown -R ${devuser}:${devuser} /home/${devuser}/Pictures
    sudo su ${devuser} --command "dbus-launch gsettings set org.gnome.desktop.background picture-uri file:///home/${devuser}/Pictures/desktop.jpg"
    sudo su ${devuser} --command "dbus-launch gsettings set org.gnome.desktop.background picture-options stretched"
fi


if [[ $(hostname) == "operator" && "${robot_has_operator}" == "0" ]];
then
    echo "######### Setting up $(hostname) machine ##############################"
    setup_group "operator"
    echo "######### Finished Setting up $(hostname) machine #####################"
elif [[ $(hostname) == "operator" && "${robot_has_operator}" == "1" ]];
then
    echo "######### Setting up $(hostname) machine ##############################"
    echo -n operator > /etc/zoox/role
    setup_group "operator_limited"
    echo "######### Finished Setting up $(hostname) machine #####################"
elif [[ $(hostname) =~ ^vision.* ]];
then
    echo "######### Setting up $(hostname) machine ##############################"
    echo -n vision > /etc/zoox/role
    setup_group "vision_limited"
    echo "######### Finished Setting up $(hostname) machine #####################"
else
    echo "######### Setting up AI machine ####################################"
    # Update the network type

    update_robot_network_type

    echo -n ai > /etc/zoox/role
    setup_group "ai"
    echo "######### Finished Setting up AI machine ###########################"

    echo "######### Setting up vision machines ###############################"
    ensure_vision_hosts_and_setup
    echo "######### Finished Setting up vision machines ######################"

    echo "######### Setting up operator ###############################"
    ensure_operator_up_and_setup
    echo "######### Finished Setting up operator ######################"

    echo "##################### Download all imageries ########################"
    download_imageries
    echo "##################### Finished Downloading all imageries ########################"

    reboot_operator
fi

if [[ ! $(hostname) =~ ^vision.* ]];
then

    # Install zstart
    sudo su ${devuser} --command "curl zstart.web.zooxlabs.com/install | sudo bash"

    if [[ ! ( $(hostname) == "operator"  && "${robot_has_operator}" == "1" ) ]];
    then
        log_message "Success! Imaging for robot $robot_name is complete" 0
    fi
fi

exit 0
