#!/usr/bin/env bash

# Copyright 2019 The Kubernetes Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#+
#+ Usage: testgridshot <release>
#+
#+   <release> can be something like 'master', '1.15', ...
#+
#+   This will inspect the testgrid dashboards 'blocking' & 'informing' for the
#+   specified <release>, create screenshots of failing tests, and print a markdown
#+   stub on standard out. This stub is intended to be copy & pasted into a issue
#+   comment for the release cutting GitHub issue.
#+
#+ Configuration:
#+   Configuration can be passed in via the environment. The following
#+   variables are available:
#+
#+   BOARDS: [default: "blocking informing"]
#+     Change which boards you are interested in
#+
#+   STATES: [default: "FAILING"]
#+     Change the tests you want to screenshot
#+     Available states: FAILING, FLAKY, PASSING
#+
#+   BUCKET: [default: "k8s-staging-releng"]
#+     The name of the bucket to upload the images to.
#+     `gsutil` is used for the upload, thus it must be installed and configured correctly.
#+     The files will be put into '/testgridshot/<release>/<datetime>_<rand>/...'
#+
#+   BLOCK_WIDTH: [default: 30]
#+     The width of the individual testgrid red & green block.
#+     Influences on how many test runs will be shown on the screenshot.
#+
#+   WIDTH/HEIGHT: [default: WIDTH=3000/HEIGHT=2500]
#+     The width and height of the generated screenshot in pixel.
#+
#+   RETRY_COUNT: [default: 3]
#+     Change how often we should retry when talking to (curl'ing) APIs.
#+
#+   RETRY_SLEEP: [default: 2]
#+     Change the amount of time we want to wait between consecutive API call
#+     retries.
#+
#+   LOCAL_IMG_DIR: [default: '']
#+     Set this to a local path to skip the upload to the GCS bucket and
#+     instead move the resulting files to the specified local path.
#+     The directory specified must not exist yet or testgridshot will fail
#+     to ensure we don't override files unintentionally.
#+
#+   Example:
#+     $ BUCKET='kubeimages' BOARDS='all informing' STATES='FAILING FLAKY' BLOCK_WIDTH=10 ./testgridshot 1.15
#+         Creates dense screenshots of all the 'sig-release-1.15-{all,informing}'
#+         boards which are either failing or flaking
#+
#+ How it works:
#+   General flow:
#+   - Get the dasbhoard URLs we are interested in from testgrid's 'summary'
#+     endpoint
#+   - Use 'render-tron.appspot.com' to create screenshots of the testgrid
#+     boards
#+   - Upload all the screenshots to a google storage bucket and make them public
#+   - Print a markdown stub which links to the images in the bucket on StdOut
#+

set -o errexit
set -o nounset
set -o pipefail

readonly BOARDS="${BOARDS:-blocking informing}"
# Available states: FAILING, FLAKY, PASSING
readonly STATES="${STATES:-FAILING}"
readonly BUCKET="${BUCKET:-k8s-staging-releng}"
readonly BLOCK_WIDTH="${BLOCK_WIDTH:-30}"
readonly WIDTH="${WIDTH:-3000}"
readonly HEIGHT="${HEIGHT:-2500}"
readonly RETRY_COUNT="${RETRY_COUNT:-3}"
readonly RETRY_SLEEP="${RETRY_SLEEP:-2}"
readonly LOCAL_IMG_DIR="${LOCAL_IMG_DIR:-}"

readonly TESTGRID='https://testgrid.k8s.io'
readonly RENDER_TRON='https://render-tron.appspot.com/screenshot'
readonly ISSUE_STUB_SUFFIX='issue_comment_part.'


get_tests_by_status() {
  local status="$1"
  shift

  local board
  local summary

  for board in "$@"
  do
    summary="$( curl_with_retry --retry 0 "${TESTGRID}/${board}/summary" 2>/dev/null )" || {
      log "Could not get summary for board '${board}'"
      return 1
    }

    jq --arg status "$status" --arg board "$board" -r '
      to_entries[]
        | select(.value.overall_status==$status)
        | $board + "#" + .key
    ' <<< "$summary"
  done
}

urlencode() {
  echo "$1" | jq -sRr @uri
}

gen_file_name() {
  echo "${1//[^a-zA-Z0-9]/_}"
}

curl_with_retry() {
  curl -fqsSL --retry "$RETRY_COUNT" --retry-delay "$RETRY_SLEEP" "$@"
}

log() {
  echo "$(get_timestamp) ${*}" >&2
}

screenshot() {
  local url="$1"
  local target="$2"
  local url_encoded rendertron_url

  url_encoded="$( urlencode "${url}" )"
  rendertron_url="${RENDER_TRON}/${url_encoded}?width=${WIDTH}&height=${HEIGHT}"

  curl_with_retry -o "${target}" "$rendertron_url"
}

get_issue_stub_name() {
  local dir="$1"
  local idx="$2"
  printf '%s/%s%05d' "$dir" "$ISSUE_STUB_SUFFIX" "$idx"
}

combine_issue_stubs() {
  local dir="$1"
  find "$dir" -mindepth 1 -maxdepth 1 -name "${ISSUE_STUB_SUFFIX}*" \
    | sort -n \
    | xargs cat
}

get_timestamp() {
  date '+%Y-%m-%d %H:%M:%S%z'
}

comma_sep() {
  echo "$*" \
    | sed -e 's@^\s\+@@' -e 's@\s\+$@@' -e 's@\s\+@, @g'
}

get_header_stub() {
  local release="$1" ; shift
  local board

  echo "### Testgrid dashboards for ${release}"

  echo "Boards checked for $(comma_sep "${STATES}"):"

  for board in "$@"
  do
    printf -- '- [%s](%s)\n' "$board" "${TESTGRID}/${board}"
  done
}

upload_and_publish_images() {
  local local_path="${1}"
  local img_bucket_dir="${2}"
  local bucket_path="gs://${BUCKET}/${img_bucket_dir}/"

  gsutil -m -q cp -r "${local_path}/." "${bucket_path}"
}

usage() {
  local usage_marker='^#\\+ ?'
  # shellcheck disable=SC2016
  local awk_prog='$0 ~ RE { gsub(RE, ""); print }'

  awk -vRE="$usage_marker" "$awk_prog" <"$0" >&2
}

wait_for_all() {
  for pid in "$@"
  do
    wait "$pid"
  done
}

shoot() {
  local target_release="$1"
  local boards
  local tests t s raw_tests
  local idx=0
  local img_bucket_dir
  local pids=()
  local comment

  if [ -n "$LOCAL_IMG_DIR" ] && [ -e "$LOCAL_IMG_DIR" ]
  then
    log "Directory '${LOCAL_IMG_DIR}' must be empty"
    return 1
  fi

  tmp_dir="$( mktemp -d )"
  trap 'rm -rf -- "$tmp_dir"' EXIT

  # that's where we locally stage the images
  local img_dir="${tmp_dir}/images"
  mkdir -p "${img_dir}"

  img_bucket_dir="testgridshot/${target_release}/$(date -u '+%Y-%m-%d_%H-%M-%S')_$(openssl rand -hex 4)"

  read -r -a boards <<< "$BOARDS"
  # prepend all elements with 'sig-release-<release>-' to form the full
  # testgrid dashborad name
  boards=( "${boards[@]/#/sig-release-${target_release}-}" )

  get_header_stub "$target_release" "${boards[@]}" > "$( get_issue_stub_name "${tmp_dir}" "${idx}" )"

  for s in ${STATES}
  do
    raw_tests="$( get_tests_by_status "$s" "${boards[@]}" )"
    mapfile -t tests <<< "$raw_tests"
    for t in "${tests[@]}"
    do
      [ -n "${t}" ] || continue

      idx=$(( idx + 1 ))
      (
        local testgrid_url timestamp file_name image_url file_base_name file_size

        local_log() {
          log "[${idx}]" "$@"
        }

        local_log "starting ${t}"

        testgrid_url="${TESTGRID}/${t}&width=${BLOCK_WIDTH}"
        file_base_name="$( gen_file_name "${t}" ).jpg"
        file_name="${img_dir}/${file_base_name}"

        # create & download screenshot
        local_log "screenshoting ${testgrid_url} to ${file_base_name}"
        screenshot "${testgrid_url}" "${file_name}"
        timestamp="$( get_timestamp )"

        read -r file_size <<< "$(wc -c < "$file_name")"
        local_log "${file_base_name}: ${file_size} bytes"

        image_url="https://storage.googleapis.com/${BUCKET}/${img_bucket_dir}/${file_base_name}"

        # generate issue section
        local_log "generating markdown stub"
        printf \
          '\n<details><summary><tt>%s</tt> %s %s <a href="%s">[testgrid]</a></summary><p>\n\n![%s](%s)\n</p></details>\n' \
          "${timestamp}" "${s}" "${t}" "${testgrid_url}" "${t}" "${image_url}" \
          > "$( get_issue_stub_name "${tmp_dir}" "$idx" )"

        local_log "done, image will be available at ${image_url}"
      )&
      pids+=( $! )
    done
  done

  wait_for_all "${pids[@]}"

  if [ "$idx" -lt 1 ]; then
    {
      echo
      echo "**NO $(comma_sep "${STATES}") TESTS**"
    } > "$( get_issue_stub_name "${tmp_dir}" $((idx + 1)) )"
  fi

  comment="$( combine_issue_stubs "${tmp_dir}" | tee "${tmp_dir}/issue_comment.txt" )"

  if [ -n "$LOCAL_IMG_DIR" ]
  then
    log "Skipping upload, moving results to '$LOCAL_IMG_DIR'"
    mv "${tmp_dir}" "${LOCAL_IMG_DIR}"
  else
    log "Starting upload to '$BUCKET'"
    upload_and_publish_images "$img_dir" "$img_bucket_dir"
    log "Upload finished successfully"
  fi

  echo

  echo '<!-- ----[ issue comment ]---- -->'
  echo "$comment"
  echo '<!-- ----[ issue comment ]---- -->'
}

main() {
  local target_release="${1:-}"

  if [ -z "$target_release" ]; then
    usage
    exit
  fi

  shoot "$target_release"
}

main "$@"
