#!/bin/bash

FIVEXX_ONLY=true
CONTAINER_SELECT=false
WRITE_TO_FILE=false
FLAGS=""
HELP_TEXT="script usage: $(basename $0) [-h (prints this message)] [-a (print all logs, not just 5xx errors)] [-c (view logs of a certain container)] [-C (view logs of all containers)] [-p (previous pod state)] [-s TIMEDELTA (print logs since TIMEDELTA)] pod_name"

while getopts 'hacCwps:' OPTION; do
  case "$OPTION" in
    h)
      echo "$HELP_TEXT" >&2 ; exit 0
      ;;
    a)
      FIVEXX_ONLY=false
      ;;
    c)
      CONTAINER_SELECT=true
      ;;
    C)
      [[ $CONTAINER_SELECT == true ]] && echo "ERROR: -C AND -c CAN'T BE USED AS PART OF THE SAME COMMAND" && exit -2
      FLAGS="$FLAGS --all-containers"
      echo "VIEWING LOGS OF ALL CONTAINERS IN POD"
      ;;
    p)
      FLAGS="$FLAGS -p"
      echo "VIEWING LOGS OF POD'S PREVIOUS STATE"
      ;;
    w)
      WRITE_TO_FILE=true
      echo "WRITE TO FILE: TRUE"
      ;;
    s)
      TIMEDELTA="$OPTARG"
      echo "TIME DELTA SPECIFIED: $TIMEDELTA"
      FLAGS="$FLAGS --since=$TIMEDELTA"
      ;;
    ?)
	echo "$HELP_TEXT" >&2 ; exit 0
      exit 0
      ;;
  esac
done
shift "$(($OPTIND -1))"

POD_ARG="$1"
ENV="$( kubectx -c | cut -d "-" -f3 )"
POD_LIST="$( kubectl get pods --all-namespaces | grep "$POD_ARG" | awk '$0!=""{print NR, $0}' )"
POD_NUM="$( echo "$POD_LIST" | wc -l )"
POD=""
CONTAINER=""

get_containers () {
    CONTAINER_LIST="$( kubectl get pods "$POD" -o jsonpath="{.spec['containers','initContainers'][*].name}" | tr ' ' '\n' | awk '$0!=""{print NR, $0}' )"
    CONTAINER_NUM="$( echo "$CONTAINER_LIST" | wc -l )"
    if [[ "$CONTAINER_NUM" -gt 1 ]]; then 
        echo -e "\nLIST OF CONTAINERS IN POD ${POD}:\n${CONTAINER_LIST}\n"
	read -p 'PLEASE SELECT POD NUMBER: ' containerindex
	[[ "$containerindex" =~ [0-9]+ ]] || echo "FUNKY INPUT, DEFAULTING TO 1" && containerindex=1
	if [[ "$containerindex" -le "$CONTAINER_NUM" ]]; then
            CONTAINER="$( echo "$CONTAINER_LIST" | sed -n "${containerindex}p" | sed -r "s|^[^ ]* *||g" )"
        else
            echo "ERROR: INDEX OF CONTAINER EXCEEDS NUMBER OF CONTAINERS FOUND INSIDE POD"
            exit 1
        fi
    elif [[ "$CONTAINER_NUM" -eq 1 ]]; then
        CONTAINER="$( echo "$CONTAINER_LIST" | sed -r "s|^[^ ]* *||g" )"
    else
        echo "NO CONTAINERS FOUND IN $POD"
        exit -1
    fi
    FLAGS="$FLAGS --container="$CONTAINER""
}

get_logs () {
    grepstr=""
    echostr="ALL LOGS FOR POD ${POD}:\n"
    tmp_file="$( mktemp -p "/tmp" "${POD}_logs_tmp_XXXXX" )"
    file="${POD}_logs.txt"
    NS="$( echo "$POD_LIST" | grep "$POD" | sed -r "s|^[^ ]* *([^ ]*) *(.*?)|\1|g" )"
    [[ ! $( kubens -c | grep "$NS" ) ]] && kubens "$NS" 
    [[   $( kubens -c | grep "$NS" ) ]] || exit -1
    [[ $CONTAINER_SELECT == true ]] && get_containers && echostr="$( echo "$echostr" | sed -r "s|:| (container $CONTAINER):|g" )"
    [[ ! $FIVEXX_ONLY ]] && echostr="$( echo "$echostr" | sed -r "s|^ALL|5XX|" )" && grepstr="[^\.0-9a-zA-Z\-]5[0-9]{2}[^\.0-9a-zA-Z\-]"
    echo -e "$echostr"
    kubectl logs "$POD" $FLAGS --timestamps=true --prefix=true | grep --color=always -E "$grepstr" | tee "$tmp_file"

    [[ "$WRITE_TO_FILE" = "true" ]] && cat "$tmp_file" | sed -e 's/\x1b\[[0-9;]*m//g' > "$file" && echo "Logs saved to '$file'."
    
    rm -f "$tmp_file"
}

[[ -z "$POD_ARG" ]] && echo "ERROR: MUST SPECIFY POD" && exit -1

echo -e "CURRENT ENV: $ENV\n"

if [[ "$POD_NUM" -gt 1 ]]; then 
    echo -e "FOUND MORE THAN ONE POD THAT MATCHES STRING:\n${POD_LIST}\n"
    read -p 'PLEASE SELECT POD NUMBER: ' podindex
    [[ "$podindex" =~ [0-9]+ ]] || echo "FUNKY INPUT, DEFAULTING TO 1" && podindex=1
    if [[ "$podindex" -le "$POD_NUM" ]]; then
        POD="$( echo "$POD_LIST" | sed -n "${podindex}p" | sed -r "s|^([^ ]* *){2}||g" | grep -o -E "^[^ ]*" )"
	get_logs
    else
	echo "ERROR: INDEX OF POD EXCEEDS NUMBER OF PODS FOUND"
        exit 1
    fi
elif [[ "$POD_NUM" -eq 1 ]]; then
    POD="$( echo "$POD_LIST" | sed -r "s|^([^ ]* *){2}||g" | grep -o -E "^[^ ]*" )"
    get_logs
else
    echo "POD NOT FOUND ON $ENV"
    exit -1
fi