#!/bin/bash
if [ -z "${KUBECTL_BIN}" ]; then
  if hash kubectl 2>/dev/null; then
    KUBECTL_BIN='kubectl'
  elif hash kubectl.exe 2>/dev/null; then
    KUBECTL_BIN='kubectl.exe'
  elif hash microk8s 2>/dev/null; then
    KUBECTL_BIN='microk8s.kubectl'
  fi
fi

if ! hash "${KUBECTL_BIN}" 2>/dev/null; then
  echo >&2 "kubectl is not installed"
  exit 1
fi

PROGNAME="$(basename "$0")"
readonly PROGNAME

calculate_default_namespace() {
  local config_namespace
  config_namespace=$(${KUBECTL_BIN} config view --minify --output 'jsonpath={..namespace}')
  echo "${KUBETAIL_NAMESPACE:-${config_namespace:-default}}"
}

# Sets default color output to 'false' if output is not a terminal
terminal_aware_default_color=line
[ ! -t 1 ] && terminal_aware_default_color=false

default_previous="${KUBETAIL_PREVIOUS:-false}"
default_since="${KUBETAIL_SINCE:-10s}"
default_namespace="$(calculate_default_namespace)"
default_follow="${KUBETAIL_FOLLOW:-true}"
default_prefix="${KUBETAIL_PREFIX:-true}"
default_line_buffered="${KUBETAIL_LINE_BUFFERED:-}"
default_colored_output="${KUBETAIL_COLORED_OUTPUT:-$terminal_aware_default_color}"
default_timestamps="${KUBETAIL_TIMESTAMPS:-}"
default_jq_selector="${KUBETAIL_JQ_SELECTOR:-}"
default_skip_colors="${KUBETAIL_SKIP_COLORS:-7,8}"
default_tail="${KUBETAIL_TAIL:--1}"
default_show_color_index="${KUBETAIL_SHOW_COLOR_INDEX:-false}"

# do not keep an unused namespace var; use namespace_arg consistently
namespace_arg="-n ${default_namespace}"
follow="${default_follow}"
prefix="${default_prefix}"
line_buffered="${default_line_buffered}"
colored_output="${default_colored_output}"
timestamps="${default_timestamps}"
jq_selector="${default_jq_selector}"
skip_colors="${default_skip_colors}"
tail="${default_tail}"
show_color_index="${default_show_color_index}"

if [[ ${1} != -* ]]; then
  pod="${1}"
fi

containers=()
selector=()
regex='substring'
previous="${default_previous}"
since="${default_since}"
version="1.6.19-SNAPSHOT"
dryrun=false
cluster=""

usage="${PROGNAME} <search term> [-h] [-c] [-n] [-t] [-l] [-f] [-d] [-P] [-p] [-s] [-b] [-e] [-j] [-k] [-z] [-v] [-r] [-i] -- tail multiple Kubernetes pod logs at the same time

where:
    -h, --help              Show this help text.
    -c, --container         The name of the container to tail in the pod (if multiple containers are defined in the pod).
                            Defaults to all containers in the pod. Can be used multiple times.
    -t, --context           The k8s context. ex. int1-context. Relies on ~/.kube/config for the contexts.
    -l, --selector          Label selector. If used the pod name is ignored.
    -n, --namespace         The Kubernetes namespace where the pods are located. Defaults to \"${default_namespace}\".
    -f, --follow            Specify if the logs should be streamed. (true|false) Defaults to ${default_follow}.
    -d, --dry-run           Print the names of the matched pods and containers, then exit.
    -P, --prefix            Specify if add the pod name prefix before each line. (true|false) Defaults to ${default_prefix}.
    -p, --previous          Return logs for the previous instances of the pods, if available. (true|false) Defaults to ${default_previous}.
    -s, --since             Only return logs newer than a relative duration like 5s, 2m, or 3h. Defaults to ${default_since}.
    -b, --line-buffered     This flags indicates to use line-buffered. (true|false) Defaults to ${default_line_buffered:-false}.
    -e, --regex             The type of name matching to use (regex|substring). Defaults to ${regex}.
    -j, --jq                If your output is json - use this jq-selector to parse it. Defaults to \"${default_jq_selector}\".
                            example: --jq \".logger + \\\" \\\" + .message\"
    -k, --colored-output    Use colored output (pod|line|false).
                            pod = only color pod name, line = color entire line, false = don't use any colors.
                            Defaults to ${default_colored_output}.
    -z, --skip-colors       Comma-separated list of colors to not use in output.
                            If you have green foreground on black, this will skip dark grey and some greens: -z 2,8,10
                            Defaults to: ${default_skip_colors}.
        --timestamps        Show timestamps for each log line. (true|false) Defaults to ${default_timestamps:-false}.
        --tail              Lines of recent log file to display. Defaults to ${default_tail}, showing all log lines.
    -v, --version           Prints the kubetail version.
    -r, --cluster           The name of the kubeconfig cluster to use.
    -i, --show-color-index  Show the color index before the pod name prefix that is shown before each log line.
                            Normally only the pod name is added as a prefix before each line, for example \"[app-5b7ff6cbcd-bjv8n]\",
                            but if \"show-color-index\" is true then color index is added as well: \"[1:app-5b7ff6cbcd-bjv8n]\".
                            This is useful if you have color blindness or if you want to know which colors to exclude (see \"--skip-colors\").
                            Defaults to ${default_show_color_index}.

examples:
    ${PROGNAME} my-pod-v1
    ${PROGNAME} my-pod-v1 -c my-container
    ${PROGNAME} my-pod-v1 -t int1-context -c my-container
    ${PROGNAME} '(service|consumer|thing)' -e regex
    ${PROGNAME} -l service=my-service
    ${PROGNAME} --selector service=my-service --since 10m
    ${PROGNAME} --tail 1"

if [ "$#" -ne 0 ]; then
  while [ "$#" -gt 0 ]; do
    case "$1" in
    -h | --help)
      echo "$usage"
      exit 0
      ;;
    -v | --version)
      echo "$version"
      exit 0
      ;;
    -c | --container)
      containers+=("$2")
      ;;
    -e | --regex)
      if [ "$2" = "substring" ]; then regex="substring"; else regex="regex"; fi
      ;;
    -t | --context)
      context="$2"
      ;;
    -r | --cluster)
      cluster="--cluster $2"
      ;;
    -l | --selector)
      selector=(--selector "$2")
      pod=""
      ;;
    -d | --dry-run)
      dryrun=true
      ;;
    -p | --previous)
      if [ "$2" = "false" ]; then previous="false"; else previous="true"; fi
      ;;
    -s | --since)
      if [ -n "$2" ]; then since="$2"; else since="${default_since}"; fi
      ;;
    -n | --namespace)
      if [ -n "$2" ]; then namespace_arg="--namespace $2"; fi
      ;;
    -f | --follow)
      if [ "$2" = "false" ]; then follow="false"; else follow="true"; fi
      ;;
    -P | --prefix)
      if [ "$2" = "false" ]; then prefix="false"; else prefix="true"; fi
      ;;
    -b | --line-buffered)
      if [ "$2" = "false" ]; then line_buffered=""; else line_buffered="| grep - --line-buffered"; fi
      ;;
    -k | --colored-output)
      if [ -n "$2" ]; then colored_output="$2"; else colored_output="${default_colored_output}"; fi
      ;;
    -j | --jq)
      if [ -n "$2" ]; then jq_selector="$2"; else jq_selector="${default_jq_selector}"; fi
      ;;
    -z | --skip-colors)
      if [ -n "$2" ]; then skip_colors="$2"; else skip_colors="${default_skip_colors}"; fi
      ;;
    --timestamps)
      if [ "$2" = "false" ]; then timestamps="$1=$2"; else timestamps="$1"; fi
      ;;
    --tail)
      if [ -n "$2" ]; then tail="$2"; else tail="${default_tail}"; fi
      ;;
    -i | --show-color-index)
      if [ -n "$2" ]; then show_color_index="$2"; else show_color_index="${default_show_color_index}"; fi
      ;;
    --)
      break
      ;;
    -*)
      echo "Invalid option '$1'. Use --help to see the valid options" >&2
      exit 1
      ;;
    *) ;;
    esac
    shift
  done
else
  echo "$usage"
  exit 1
fi

# Join function that supports a multi-character separator
join() {
  local sep="$1"
  shift
  local F=0 x
  for x in "$@"; do
    if [[ $F -eq 1 ]]; then
      echo -n "$sep"
    else
      F=1
    fi
    echo -n "$x"
  done
  echo
}

# If comma-separated pod names provided (without --regex), convert to regex
if [[ ${pod} == *","* ]] && [ ! "${regex}" = 'regex' ]; then
  IFS=',' read -r -a pods_to_match <<<"${pod}"
  pod="$(join '.*|.*' "${pods_to_match[@]}")"
  pod=".*${pod}.*"
  regex='regex'
fi

grep_matcher=''
if [ "${regex}" = 'regex' ]; then
  echo "Using regex '${pod}' to match pods"
  grep_matcher='-E'
fi

# Build grep args safely
grep_args=(--color=never)
if [ -n "${grep_matcher}" ]; then grep_args+=("${grep_matcher}"); fi

# Get matching pods safely into an array
mapfile -t matching_pods < <(
  ${KUBECTL_BIN} get pods \
    ${context:+--context="${context}"} \
    "${selector[@]}" \
    "${namespace_arg}" "${cluster}" \
    --field-selector=status.phase=Running \
    --output=jsonpath='{.items[*].metadata.name}' |
    xargs -n1 |
    grep "${grep_args[@]}" -- "${pod}"
)

matching_pods_size=${#matching_pods[@]}
if [ "${matching_pods_size}" -eq 0 ]; then
  echo "No pod exists that matches ${pod}"
  exit 1
fi

color_end="$(tput sgr0)"

display_names_preview=()
logs_commands=() # was pod_logs_commands, but that var was unused
i=0
color_index=0

next_col() {
  local potential_col=$(($1 + 1))
  if [[ $skip_colors =~ (^|,)$potential_col($|,) ]]; then
    next_col "$potential_col"
  else
    echo "$potential_col"
  fi
}

# Allows for more colors, useful if tailing many pods
if [ "${colored_output}" != "false" ]; then
  export TERM=xterm-256color
fi

kill_kubectl_processes() { kill 0; }
trap kill_kubectl_processes EXIT

# Cache pod/container names (avoid multiple API calls)
all_pods_containers="$(
  ${KUBECTL_BIN} get pods \
    "${namespace_arg}" \
    ${context:+--context="${context}"} \
    --output=jsonpath='{range .items[*]}{.metadata.name} {.spec['"'"'containers'"'"', '"'"'initContainers'"'"'][*].name} \n{end}'
)"

for pod in "${matching_pods[@]}"; do
  if [ ${#containers[@]} -eq 0 ]; then
    readarray -t pod_containers < <(
      printf '%s\n' "${all_pods_containers}" |
        grep -F -- "${pod} " |
        cut -d ' ' -f2- |
        xargs -n1
    )
  else
    pod_containers=("${containers[@]}")
  fi

  for container in "${pod_containers[@]}"; do
    if [[ ${matching_pods_size} -eq 1 && ${#pod_containers[@]} -eq 1 ]]; then
      single_stream="true"
    else
      single_stream="false"
    fi

    if [ "${colored_output}" = "false" ] || [ "${single_stream}" = "true" ]; then
      color_start="$(tput sgr0)"
      color_index_prefix=""
    else
      color_index="$(next_col "${color_index}")"
      color_start="$(tput setaf "${color_index}")"
      if [ "${show_color_index}" = "true" ]; then
        color_index_prefix="${color_index}:"
      else
        color_index_prefix=""
      fi
    fi

    if [ ${#pod_containers[@]} -eq 1 ]; then
      display_name="${pod}"
    else
      display_name="${pod} ${container}"
    fi

    if [ "${colored_output}" = "false" ]; then
      display_names_preview+=("${display_name}")
    else
      display_names_preview+=("${color_index_prefix}${color_start}${display_name}${color_end}")
    fi

    if [ "${prefix}" = "false" ]; then
      prefix_line=""
    else
      if [ "${colored_output}" = "false" ]; then
        prefix_line="[${display_name}] "
      else
        prefix_line="${color_start}[${color_end}${color_index_prefix}${color_start}${display_name}]${color_end} "
      fi
    fi

    if [ "${colored_output}" = "false" ] || [ "${colored_output}" = "pod" ]; then
      colored_line="${prefix_line}\$REPLY"
    else
      colored_line="${prefix_line}${color_start}\$REPLY${color_end}"
    fi

    kubectl_cmd="${KUBECTL_BIN} ${context:+--context=${context}} logs ${pod} ${container} -f=${follow} --previous=${previous} --since=${since} --tail=${tail} ${namespace_arg} ${cluster}"
    colorify_lines_cmd='while read -r; do echo "'"${colored_line}"'" | tail -n +1; done'

    if [ -z "${jq_selector}" ]; then
      logs_commands+=("${kubectl_cmd} ${timestamps} | ${colorify_lines_cmd}")
    else
      logs_commands+=("${kubectl_cmd} | jq --unbuffered -r -R --stream '. as \$line | try (fromjson | ${jq_selector}) catch \$line' | ${colorify_lines_cmd}")
    fi

    i=$(((i + 1) % 13))
  done
done

echo "Will tail ${#display_names_preview[@]} logs..."
for preview in "${display_names_preview[@]}"; do
  echo "$preview"
done

if [[ ${dryrun} == true ]]; then
  exit 0
fi

command_to_tail="$(join " & " "${logs_commands[@]}")"

tail_follow_command="-f"
if [[ ${follow} == false ]]; then
  tail_follow_command=""
fi

# shellcheck disable=SC2086
tail ${tail_follow_command} -n +1 <(eval "${command_to_tail}") "$line_buffered"
