import argparse
import json
import os
import sys
from datetime import datetime


def format_bytes(size):
    # Define the units and their corresponding sizes
    units = ["B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]

    # Loop through each unit
    for unit in units:
        if size < 1024:
            return f"{size:.2f} {unit}"
        size /= 1024

    return f"{size:.2f} YB"


def lines(fd):
    for line in map(lambda l: l.strip(), fd):
        yield line


def main():
    parser = argparse.ArgumentParser(description="Makes statistics about kunai logs")
    parser.add_argument(
        "-r",
        "--refresh",
        default=300,
        type=int,
        help="Refresh rate in numbers of events",
    )
    parser.add_argument(
        "kunai_json_input", help="Input file in json line format or stdin with -"
    )

    args = parser.parse_args()

    if args.kunai_json_input == "-":
        lg = lines(sys.stdin)
    else:
        lg = lines(open(args.kunai_json_input, "r"))

    stats = {}
    count, size = 0, 0
    first, last = 0, 0
    for line in lg:
        event = json.loads(line)
        timestamp = datetime.fromisoformat(event["info"]["utc_time"])
        eid = event["info"]["event"]["id"]
        ename = event["info"]["event"]["name"]
        key = (eid, ename)
        if key not in stats:
            stats[key] = 0
        stats[key] += 1
        count += 1
        size += len(line)

        if first == 0:
            first = timestamp
        last = timestamp

        if count != 0 and count % args.refresh == 0:
            os.system("clear")
            delta = last - first

            if delta.total_seconds() == 0:
                continue

            for k in sorted(stats, key=lambda k: k[0]):
                c = stats[k]
                eps = c / delta.total_seconds()
                print(f"{k[1]}:".ljust(20), end=" ")
                print(f"{c} ->".rjust(10), end=" ")
                print(f"{eps:.2f} e/s".rjust(12))
            tot_eps = count / delta.total_seconds()
            tot_bps = size / delta.total_seconds()
            print()
            print("Totals")
            print("Events:".ljust(20), end=" ")
            print(f"{count} ->".rjust(10), end=" ")
            print(f"{tot_eps:.2f} e/s".rjust(12))
            # size in bytes
            print("Bytes:".ljust(20), end=" ")
            print(f"{format_bytes(size)} ->".rjust(10), end=" ")
            print(f"{format_bytes(tot_bps)}/s".rjust(12))
            # time
            print("Sample Duration:".ljust(20), end=" ")
            print(f"{delta}".rjust(10))


if __name__ == "__main__":
    main()
