#!/bin/sh

# Configure adb to accept adb connections via TCP/IP and make sure that the device is actually
# accessible.
# This function assumes that the device has a network address. Guards around `adb tcpip` and test
# connection setups using `adb connect` are used to check if the device is reachable after this
# call.
# Globals:
#   dut_address         Set to "ip_address:adb_port" by this function
# Arguments:
#   adb_tcpip_attempts  Number of tries for enabling adb TCP/IP mode on the device
#   timeout_secs        Timeout for waiting for getting the IP address from the device
#   adb_port            Network port to use for adb TCP/IP
# Returns:
#   0 only if the device is accessible via adb TCP/IP, 1 otherwise.
open_adb_tcpip_on_local_device() {
    [ "$#" -lt 2 ] || [ "$#" -gt 3 ] && \
        error_fatal "Usage: open_adb_tcpip_on_local_device adb_tcpip_attempts timeout_secs [adb_port]"
    # shellcheck disable=SC2039
    local adb_tcpip_attempts="$1"
    # shellcheck disable=SC2039
    local timeout_secs="$2"
    # shellcheck disable=SC2039
    local adb_port="$3"
    if [ -z "${adb_port}" ]; then
        # shellcheck disable=SC2039
        local adb_port=5555    # default port assumed by adb connect
    fi

    # shellcheck disable=SC2039
    local end=$(( $(date +%s) + timeout_secs ))

    # shellcheck disable=SC2039
    local ret_val=0

    # adb tcpip may fail for different reasons
    # (e.g., "error: protocol fault (couldn't read status): Connection reset by peer").
    # Just hope that it works after a few retries.
    # shellcheck disable=SC2039
    local adb_tcpip_retry_wait_secs=10
    # shellcheck disable=SC2039
    local attempt=0
    while [ "${attempt}" -lt "${adb_tcpip_attempts}" ] && [ "$(date +%s)" -lt "$end" ]; do
        ret_val=0
        adb tcpip "${adb_port}" || ret_val=$?
        if [ "${ret_val}" -eq 0 ]; then
            break
        fi
        info_msg "adb tcpip apparently failed. Retrying in a moment..."
        sleep "${adb_tcpip_retry_wait_secs}"
        adb usb || true # In between, make sure to have some default state.
    done

    if [ "${ret_val}" -ne 0 ]; then
        warn_msg "Could not prepare the device for adb TCP/IP connections: adb tcpip failed."
        return 1
    fi

    # `adb tcpip` sometimes takes some time
    # (on some builds, up to 10 seconds were observed)
    # shellcheck disable=SC2039
    local success=false
    while [ "$(date +%s)" -lt "$end" ]; do
        # Detect the DUT IP address inside the loop. That fixes cases where the IP address changes
        # soon after boot (and thus differs between the `get_ip_address` and `adb connect` calls)
        # but is static from there on.
        # shellcheck disable=SC2039
        local ip_address
        ret_val=0
        ip_address="$(get_ip_address "${timeout_secs}")" || ret_val=$?
        if [ "${ret_val}" -ne 0 ] || [ -z "${ip_address}" ]; then
            warn_msg "get_ip_address failed unexpectedly. Retrying in 5 seconds..."
            sleep 5s
            continue
        fi
        dut_address="${ip_address}:${adb_port}"
        if [ "$(adb connect "${dut_address}" | grep -c '^connected to ')" -eq 1 ]; then
            success=true
            break
        fi
        sleep 1
    done

    # Make sure the device is not reserved to the local adb server.
    adb disconnect "${dut_address}" >/dev/null 2>&1 || true

    if [ "${success}" = false ]; then
        warn_msg "Could not prepare the device for adb TCP/IP connections: device is not reachable via network."
        return 1
    fi
}

# Make this device accessible via adb TCP/IP and send its address via handshake to a waiting role.
# NOTE: This function must only be called once per role per test submission, as LAVA does not allow
# to use the same MultiNode message ID multiple times.
# One job instance must call connect_to_remote_adb_tcpip_devices to receive the send addresses and
# complete the handshake.
# See open_adb_tcpip_on_local_device and connect_to_remote_adb_tcpip_devices
# Globals:
#   dut_address         Set to "ip_address:adb_port" by this function
# Arguments:
#   adb_tcpip_attempts  Number of tries for establishing enabling adb TCP/IP mode on the device
#   timeout_secs        Timeout for waiting for getting the IP address from the device
#   adb_port            Network port to use for adb TCP/IP
# Returns:
#   0 only if the device is accessible via adb TCP/IP, 1 otherwise.
share_local_device_over_adb_tcpip() {
    [ "$#" -lt 2 ] || [ "$#" -gt 3 ] && \
        error_fatal "Usage: share_local_device_over_adb_tcpip adb_tcpip_attempts timeout_secs [adb_port]"
    # shellcheck disable=SC2039
    local adb_tcpip_attempts="$1"
    # shellcheck disable=SC2039
    local timeout_secs="$2"
    # shellcheck disable=SC2039
    local adb_port="$3"

    # shellcheck disable=SC2039
    local ret_val=0
    open_adb_tcpip_on_local_device "${adb_tcpip_attempts}" "${timeout_secs}" "${adb_port}" || ret_val=$?
    if [ "${ret_val}" -ne 0 ]; then
        return "${ret_val}"
    fi

    lava-sync start_handover
    lava-send dut_address dut_address="${dut_address}"
    lava-sync finish_handover
}

# Counterpart to share_local_device_over_adb_tcpip
# Wait for other job instances to send their device address, guarded by a handshake for
# synchronization.
# Globals:
#   None
# Arguments:
#   adb_connect_timeout_secs    Timeout for waiting for getting the IP address from the device
#   device_worker_mapping_file  File to store a mapping between devices and their LAVA worker host
#       in the format 'serial_or_address;worker_host_id'. This file is relevant for following
#       functions to communicate with the devices or workers.
#       Optional: This file will not be created if no path is specified.
# Returns:
#   0 on success 1 otherwise.
connect_to_remote_adb_tcpip_devices() {
    [ "$#" -lt 1 ] || [ "$#" -gt 2 ] && \
        error_fatal "Usage: connect_to_remote_adb_tcpip_devices adb_connect_timeout_secs [device_worker_mapping_file]"

    # shellcheck disable=SC2039
    local adb_connect_timeout_secs="$1"
    # shellcheck disable=SC2039
    local device_worker_mapping_file="$2"

    lava-sync start_handover
    # For lava-wait-all, all involved nodes must invoke lava-send with the same message id,
    # otherwise lava-wait-all would lead to a dead lock.
    # However, only the nodes that make their device accessible (workers) add the value
    # dut_address="address:port".
    lava-send dut_address
    lava-wait-all dut_address

    # The MultiNode cache file might not exist if there is no other worker.
    # shellcheck disable=SC2039
    local cache_lines
    # shellcheck disable=SC2039
    local device_worker_mapping=""
    if [ -f "/tmp/lava_multi_node_cache.txt" ]; then
        cache_lines="$(grep "dut_address" "/tmp/lava_multi_node_cache.txt" | grep -v '^$' || true)"

        for line in ${cache_lines}; do
            # <worker_job_id>:dut_address=<dut_address>
            # shellcheck disable=SC2039
            local dut_address
            dut_address="$(echo "$line" | sed 's/.*dut_address=//')"
            # shellcheck disable=SC2039
            local worker_host=
            worker_host="$(echo "$line" | cut -d: -f1)"
            device_worker_mapping="${device_worker_mapping}${dut_address};${worker_host}\n"
        done
        device_worker_mapping="$(printf "%b" "${device_worker_mapping}")"
    fi

    lava-sync finish_handover
    # adb is not super reliable, it too often sees connected and authorized devices as "offline"
    adb kill-server || true

    # Connect to remote devices and wait until they appear online

    for device_to_worker in ${device_worker_mapping}; do
        # shellcheck disable=SC2039
        local device=
        device="$(echo "${device_to_worker}" | cut -d';' -f1)"
        for _ in $(seq 5); do
            # shellcheck disable=SC2039
            local ret_val=0
            adb connect "${device}" || ret_val="$?"
            if [ "${ret_val}" -eq 0 ]; then
                break
            fi
            adb disconnect "${device}" || true
            warn_msg "adb connect failed. Retrying in a minute..."
            sleep 1m
        done
    done

    for device_to_worker in ${device_worker_mapping}; do
        # shellcheck disable=SC2039
        local device
        device="$(echo "${device_to_worker}" | cut -d';' -f1)"
        if ! timeout "${adb_connect_timeout_secs}" adb -s "${device}" wait-for-device; then
            warn_msg "adb wait-for-device for ${device} timed out after ${adb_connect_timeout_secs} seconds."
            return 1
        fi
    done

    # shellcheck disable=SC2039
    local num_remote_devices
    num_remote_devices="$(echo "${device_worker_mapping}" | wc -l)"
    info_msg "All ${num_remote_devices} remote devices are connected and online."

    info_msg "Now adding devices locally connected via USB."
    for _ in $(seq 5); do
        # shellcheck disable=SC2039
        local connected_devices
        connected_devices="$(adb devices | grep -E '^([:\.[:alnum:]]+)\s+device$' | cut -f1)"
        if [ -z "${connected_devices}" ]; then
            warn_msg "\"adb devices\" did not list any devices. Retrying in a minute..."
            sleep 1m
        fi
    done
    if [ -z "${connected_devices}" ]; then
        lava-test-raise "\"adb devices\" did not list any devices."
    fi
    # shellcheck disable=SC2039
    local remote_only_mapping="${device_worker_mapping}"
    device_worker_mapping="${device_worker_mapping}\n"
    for device in ${connected_devices}; do
        if [ "$(echo "${remote_only_mapping}" | cut -d';' -f1 | grep -xc "${device}")" -eq 0 ]; then
            device_worker_mapping="${device_worker_mapping}${device};\n"
            info_msg "Local device: ${device}"
        fi
    done
    device_worker_mapping="$(printf "%b" "${device_worker_mapping}")"

    if [ "${device_worker_mapping_file}" ]; then
        # Make mapping between attached DUTs and worker job ids and accessible to subsequent tests:
        echo "${device_worker_mapping}" > "${device_worker_mapping_file}"
        info_msg "Mapping between devices and to worker job ids stored in ${device_worker_mapping_file}:"
        info_msg "${device_worker_mapping}"
    else
        info_msg "NOT storing device to worker job id mapping, empty filename specified."
    fi
}
