#!/usr/bin/env bash

set -uo pipefail

readonly NAMESPACE="tidb-serverless"
readonly POLL_SECONDS=10
readonly DISCOVERY_TIMEOUT_SECONDS=30
readonly MAX_NOTIFICATION_CHARS=16000

usage() {
  cat <<'EOF'
Usage:
  watch-instance-jobs.sh INSTANCE

Watch Jobs labeled app.kubernetes.io/instance=INSTANCE after they have been
applied. Successful completion and unexpected states are sent directly to the
Lark webhook in LARK_BOT_WEBHOOK.

Environment:
  LARK_BOT_WEBHOOK          Required Lark custom-bot webhook URL
  JOB_WATCH_TIMEOUT_SECONDS Optional overall timeout; 0 disables it (default)
EOF
}

die() {
  echo "error: $*" >&2
  exit 1
}

if (($# != 1)); then
  usage >&2
  exit 1
fi

instance="$1"
if ((${#instance} > 63)) || [[ ! "$instance" =~ ^[A-Za-z0-9]([A-Za-z0-9_.-]*[A-Za-z0-9])?$ ]]; then
  die "invalid Kubernetes instance label value '$instance'"
fi

for required_command in kubectl jq curl; do
  command -v "$required_command" >/dev/null || die "$required_command is required"
done

webhook_url="${LARK_BOT_WEBHOOK:-}"
[[ -n "$webhook_url" ]] || die "LARK_BOT_WEBHOOK is required"

watch_timeout="${JOB_WATCH_TIMEOUT_SECONDS:-0}"
[[ "$watch_timeout" =~ ^[0-9]+$ ]] || die "JOB_WATCH_TIMEOUT_SECONDS must be a non-negative integer"

if ! kubectl_context="$(kubectl config current-context)" || [[ -z "$kubectl_context" ]]; then
  die "cannot determine the current kubectl context"
fi

selector="app.kubernetes.io/instance=${instance}"
error_file="$(mktemp)" || die "cannot create temporary file"
trap 'rm -f -- "$error_file"' EXIT

k() {
  kubectl --context "$kubectl_context" --request-timeout=15s \
    -n "$NAMESPACE" "$@"
}

jobs_json='{"items":[]}'
query_error=""

get_jobs() {
  local attempt output

  for attempt in 1 2 3; do
    : >"$error_file"
    if output="$(k get jobs -l "$selector" -o json 2>"$error_file")" \
      && jq -e '.items | type == "array"' >/dev/null 2>&1 <<<"$output"; then
      jobs_json="$output"
      query_error=""
      return 0
    fi
    query_error="$(<"$error_file")"
    [[ -n "$query_error" ]] || query_error="kubectl returned invalid Job JSON"
    ((attempt < 3)) && sleep 2
  done
  return 1
}

send_notification() {
  local title="$1"
  local reason="$2"
  local conditions jobs_table pods_table failed_jobs job_name details=""
  local payload response message command_output

  conditions="$(
    jq -r '
      if (.items | length) == 0 then "(no matching Jobs)" else
        .items[]
        | ([.status.conditions[]? | select(.status == "True")
            | "\(.type)(\(.reason // "-"): \((.message // "-") | gsub("[\\r\\n]+"; " ")))"
          ] | join("; ")) as $conditions
        | "job=\(.metadata.name) app=\(.metadata.labels["app.kubernetes.io/name"] // "-") "
          + "active=\(.status.active // 0) succeeded=\(.status.succeeded // 0) "
          + "failed=\(.status.failed // 0) conditions=\(if $conditions == "" then "-" else $conditions end)"
      end
    ' <<<"$jobs_json"
  )"
  jobs_table="$(k get jobs -l "$selector" -o wide 2>&1)"
  pods_table="$(k get pods -l "$selector" -o wide 2>&1)"

  failed_jobs="$(
    jq -r '
      .items[]
      | select([.status.conditions[]?
          | select((.type == "Failed" or .type == "FailureTarget") and .status == "True")
        ] | length > 0)
      | .metadata.name
    ' <<<"$jobs_json"
  )"
  while IFS= read -r job_name; do
    [[ -n "$job_name" ]] || continue
    command_output="$(
      k get events --field-selector "involvedObject.name=$job_name" \
        --sort-by=.lastTimestamp 2>&1
    )"
    details+=$'\n\n'"Events for job/$job_name:"$'\n'"$command_output"
    command_output="$(
      k logs "job/$job_name" --all-containers=true --prefix=true --tail=80 2>&1
    )"
    details+=$'\n\n'"Logs for job/$job_name (last 80 lines):"$'\n'"$command_output"
  done <<<"$failed_jobs"

  message="$(
    printf '%s\n' \
      "$title" \
      "time: $(date '+%Y-%m-%dT%H:%M:%S%z')" \
      "kubectl context: $kubectl_context" \
      "namespace: $NAMESPACE" \
      "instance: $instance" \
      "reason: $reason" \
      "" "Job conditions:" "$conditions" \
      "" "Jobs:" "$jobs_table" \
      "" "Pods:" "$pods_table" \
      "$details"
  )"

  payload="$(
    jq -cn --arg text "$message" --argjson limit "$MAX_NOTIFICATION_CHARS" '
      ($text | if length > $limit then .[0:$limit] + "\n...[truncated]" else . end) as $body
      | {msg_type: "text", content: {text: $body}}
    '
  )"
  if ! response="$(
    curl --silent --show-error --fail --retry 3 --connect-timeout 10 --max-time 30 \
      -H 'Content-Type: application/json' --data-binary "$payload" "$webhook_url"
  )"; then
    echo "error: webhook request failed" >&2
    echo "$message" >&2
    return 1
  fi
  if ! jq -e '
    if has("StatusCode") then .StatusCode == 0
    elif has("code") then .code == 0
    else false end
  ' >/dev/null 2>&1 <<<"$response"; then
    echo "error: webhook rejected notification: $response" >&2
    echo "$message" >&2
    return 1
  fi
}

alert_and_exit() {
  send_notification "Kubernetes Job watcher detected an unexpected condition." "$1" || true
  exit 1
}

discovery_deadline=$(($(date +%s) + DISCOVERY_TIMEOUT_SECONDS))
while :; do
  get_jobs || alert_and_exit "kubectl query failed: $query_error"
  job_count="$(jq '.items | length' <<<"$jobs_json")"
  ((job_count > 0)) && break
  (($(date +%s) >= discovery_deadline)) \
    && alert_and_exit "no matching Jobs appeared within ${DISCOVERY_TIMEOUT_SECONDS}s"
  sleep "$POLL_SECONDS"
done

expected_jobs="$(jq -r '.items[] | "\(.metadata.name)\t\(.metadata.uid)"' <<<"$jobs_json" | LC_ALL=C sort)"
expected_count="$(jq '.items | length' <<<"$jobs_json")"
started_at=$(date +%s)

echo "Watching $expected_count Job(s) for instance '$instance' in context '$kubectl_context'."

while :; do
  get_jobs || alert_and_exit "kubectl query failed: $query_error"

  current_jobs="$(jq -r '.items[] | "\(.metadata.name)\t\(.metadata.uid)"' <<<"$jobs_json" | LC_ALL=C sort)"
  [[ "$current_jobs" == "$expected_jobs" ]] \
    || alert_and_exit "matching Job set changed after monitoring began"

  failed_count="$(
    jq '[.items[] | select([.status.conditions[]?
      | select((.type == "Failed" or .type == "FailureTarget") and .status == "True")
    ] | length > 0)] | length' <<<"$jobs_json"
  )"
  ((failed_count == 0)) || alert_and_exit "$failed_count Job(s) entered a failed condition"

  completed_count="$(
    jq '[.items[] | select([.status.conditions[]?
      | select(.type == "Complete" and .status == "True")
    ] | length > 0)] | length' <<<"$jobs_json"
  )"
  if ((completed_count == expected_count)); then
    echo "All $expected_count Job(s) completed successfully."
    send_notification \
      "Kubernetes Job watcher completed successfully." \
      "all $expected_count Job(s) completed successfully" || exit 1
    exit 0
  fi

  if ((watch_timeout > 0)) && (($(date +%s) - started_at >= watch_timeout)); then
    alert_and_exit "monitor timeout after ${watch_timeout}s"
  fi
  sleep "$POLL_SECONDS"
done