#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Author: Chmouel Boudjnah <chmouel@chmouel.com>
#
# 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.
import argparse
import json
import os
import re
import subprocess
import sys

failed_containers = (
    "ImagePullBackOff",
    "CrashLoopBackOff",
)


def colourText(text, color):
    colours = {
        "red": "\033[0;31m",
        "yellow": "\033[0;33m",
        "blue": "\033[0;34m",
        "cyan": "\033[0;36m",
        "cyan_italic": "\033[3;37m",
        "green": "\033[0;32m",
        "grey": "\033[0;30m",
        "magenta": "\033[0;35m",
        "white": "\033[0;37m",
        "white_bold": "\033[1;37m",
        "reset": "\033[0;0m",
    }
    s = f"{colours[color]}{text}{colours['reset']}"
    return s


def show_log(kctl, args, container, pod):
    cmd = "%s logs --tail=%s %s -c%s" % (kctl, args.maxlines, pod, container)
    lastlog = subprocess.run(
        cmd.split(" "), stderr=subprocess.PIPE, stdout=subprocess.PIPE
    )
    if lastlog.returncode != 0:
        print("i could not run '%s'" % (cmd))
        sys.exit(1)
    return lastlog.stdout.decode().strip()


def overcnt(jeez, kctl, pod, args):
    for container in jeez:
        errmsg = ""
        if args.restrict:
            if len(re.findall(args.restrict, container["name"])) == 0:
                continue

        state = list(container["state"].keys())[0].capitalize()
        if state in "Running":
            state = colourText(state, "blue")
        elif state == "Terminated":
            if container["state"]["terminated"]["exitCode"] != 0:
                state = colourText("FAIL", "red")
            else:
                state = colourText("SUCCESS", "green")
        elif state == "Waiting":
            reason = container["state"]["waiting"]["reason"]
            if reason in failed_containers:
                state = colourText(reason, "red")
                if (
                    "lastState" in container
                    and "terminated" in container["lastState"]
                    and "message" in container["lastState"]["terminated"]
                ):
                    errmsg = container["lastState"]["terminated"]["message"]
                elif (
                    "waiting" in container["state"]
                    and "message" in container["state"]["waiting"]
                ):
                    errmsg = container["state"]["waiting"]["message"]
                else:
                    errmsg = ""
            else:
                state = colourText(state + "    " + reason, "yellow")

        cname = colourText(container["name"], "white")

        line_new = " {:60}  {:>20}".format(cname, state)
        print(line_new)

        if errmsg:
            print()
            print(colourText("Error message for container %s:" % (cname), "cyan"))
            print("\n".join(["  " + i for i in errmsg.split("\n")]))
        if args.showlog:
            outputlog = show_log(kctl, args, container["name"], pod)
            if outputlog:
                print()
                print(colourText(" Logs for %s:" % (cname), "cyan"))
                print("\n".join(["  " + i for i in outputlog.split("\n")]))
                print()


def lensc(jeez):
    s = 0
    for i in jeez:
        if (
            "waiting" in i["state"]
            and i["state"]["waiting"]["reason"] in failed_containers
        ):
            s += 1
        if "terminated" in i["state"] and i["state"]["terminated"]["exitCode"] == 0:
            s += 1
    return s


def hasfailure(jeez):
    for i in jeez:
        if (
            "waiting" in i["state"]
            and i["state"]["waiting"]["reason"] in failed_containers
        ):
            return True
        if "terminated" in i["state"] and i["state"]["terminated"]["exitCode"] != 0:
            return True
    return False


def getstatus(hasfailures, allc, allf):
    if hasfailures:
        colour = "red"
        text = "FAIL"
    elif allc != allf:
        colour = "blue"
        text = "RUNNING"
    else:
        colour = "green"
        text = "SUCCESS"
    return (colour, text)


def which(program):
    import os

    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, _ = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None


def print_labels_annotations(jeez, key, label):
    labels = [
        f"  {colourText(v[0], 'white')}: {v[1]}"
        for v in jeez["metadata"][key].items()
        if v
    ]
    print(f"{colourText(label, 'cyan')}: ")
    print("\n".join(labels))
    print()


def main(args):
    kctl = "kubectl"
    if args.namespace:
        kctl += f" -n {args.namespace}"

    myself = which("kss")
    preview = f"{kctl} describe {{}}"
    if myself:
        preview = f"{myself}"
        preview += f" { '-n ' + args.namespace if args.namespace else ''} {'-A' if args.annotations else ''} {'-L' if args.labels else ''}"
        preview += " {}"

    query_args = ""
    if args.pod:
        query_args = f"-q '{' '.join(args.pod)}'"
    runcmd = f"{kctl} get pods -o name|fzf -0 -n 1 -m -1 {query_args} --preview-window 'down,75%:nowrap' --preview='{preview}'"
    args.pod = os.popen(runcmd).read().strip().replace("pod/", "").split("\n")

    if not args.pod or not args.pod[0]:
        print("No pods is no news which is arguably no worries.")
        sys.exit(1)

    for pod in args.pod:
        if not pod.strip():
            continue
        cmdline = f"{kctl} get pod {pod} -ojson"
        shell = subprocess.run(
            # "cat /tmp/a.json".split(" "),
            cmdline.split(" "),
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        if shell.returncode != 0:
            print("The was some problem running '%s'" % (cmdline))
            sys.exit(1)

        output = shell.stdout.decode().strip()
        jeez = json.loads(output)

        if "initContainerStatuses" not in jeez["status"]:
            jeez["status"]["initContainerStatuses"] = {}

        cnt_failicontainers = lensc(jeez["status"]["initContainerStatuses"])
        cnt_allicontainers = len(jeez["status"]["initContainerStatuses"])
        cnt_failcontainers = lensc(jeez["status"]["containerStatuses"])
        cnt_allcontainers = len(jeez["status"]["containerStatuses"])

        header = f"{colourText('Pod', 'cyan')}: {pod}\n"
        header += f"{colourText('Status', 'cyan')}: "

        colour, text = getstatus(
            hasfailure(jeez["status"]["initContainerStatuses"])
            or hasfailure(jeez["status"]["containerStatuses"]),
            cnt_allcontainers + cnt_allicontainers,
            cnt_failcontainers + cnt_failicontainers,
        )
        header += f"{colourText(text, colour)}"

        print(header)

        if args.labels:
            print_labels_annotations(jeez, "labels", "Labels")
        if args.annotations:
            print_labels_annotations(jeez, "annotations", "Annotations")

        if jeez["status"]["initContainerStatuses"]:
            colour, _ = getstatus(
                hasfailure(jeez["status"]["initContainerStatuses"]),
                cnt_allicontainers,
                cnt_failicontainers,
            )
            s = f"{cnt_failicontainers}/{cnt_allicontainers}"
            print(f"{colourText('Init Containers', 'cyan')}: {colourText(s, colour)}")
            overcnt(jeez["status"]["initContainerStatuses"], kctl, pod, args)
            print()

        colour, text = getstatus(
            hasfailure(jeez["status"]["containerStatuses"]),
            cnt_allcontainers,
            cnt_failcontainers,
        )
        if text == "RUNNING":
            s = cnt_allcontainers
        else:
            s = f"{cnt_failcontainers}/{cnt_allcontainers}"
        print(f"{colourText('Containers', 'cyan')}: {colourText(s, colour)}")
        overcnt(jeez["status"]["containerStatuses"], kctl, pod, args)

        if args.events:
            print(colourText("Events", "cyan"))
            cmd = f"kubectl get events --field-selector involvedObject.name={pod} --field-selector involvedObject.kind=Pod"
            output = os.popen(cmd).read().strip()
            output = "\n".join(["  " + i for i in output.split("\n")])
            output = "\n".join(
                [colourText(output.split("\n")[0], "white_bold")]
                + output.split("\n")[1:]
            )
            print(output)
        if len(args.pod) > 1:
            print()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("pod", nargs="*", default="")
    parser.add_argument("-n", "--namespace", dest="namespace", type=str)
    parser.add_argument(
        "-r",
        "--restrict",
        type=str,
        help="Restrict to show only those containers (regexp)",
    )

    parser.add_argument(
        "--events",
        "-E",
        action="store_true",
        default=False,
        help="Show events",
    )

    parser.add_argument(
        "--labels",
        "-L",
        action="store_true",
        default=False,
        help="Show labels",
    )

    parser.add_argument(
        "--annotations",
        "-A",
        action="store_true",
        default=False,
        help="Show Annotations",
    )

    parser.add_argument(
        "-l",
        "--showlog",
        action="store_true",
        default=False,
        help="Show logs of containers",
    )
    parser.add_argument(
        "--maxlines", type=str, default="-1", help="Maximum line when showing logs"
    )

    main(parser.parse_args(sys.argv[1:]))