#!/bin/env python3

import argparse
import json

import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots


def calculate_CPU_usage(record):
    """Return the used CPU delta (in ms) and CPU usage in percent since last tick"""
    cpu_percent = 0.0
    cpu_delta = (
        record["cpu_stats"]["cpu_usage"]["total_usage"]
        - record["precpu_stats"]["cpu_usage"]["total_usage"]
    )
    system_delta = (
        record["cpu_stats"]["system_cpu_usage"]
        - record["precpu_stats"]["system_cpu_usage"]
    )
    cpu_count = record["cpu_stats"]["online_cpus"]
    if system_delta > 0 and cpu_delta > 0:
        cpu_percent = (cpu_delta / system_delta) * cpu_count * 100
    if cpu_delta > 0:
        cpu_delta = cpu_delta / 1000000

    return (cpu_delta, cpu_percent)


def read_json(fileobj):
    j = []
    first_line = True
    for record in fileobj:
        # Skip the first line it does not have previous recordings
        if first_line:
            first_line = False
            continue
        try:
            record = json.loads(record)
        except json.decoder.JSONDecodeError:
            # probably last line not completely flushed
            continue
        cpu_delta, cpu_percent = calculate_CPU_usage(record)

        tick = dict(read=record["read"])
        tick["cpu_delta"] = cpu_delta
        tick["cpu_percent"] = cpu_percent
        # Calculate the time (in ms) the container was throttled since last tick
        tick["cpu_throttled_ms"] = (
            record["cpu_stats"]["throttling_data"]["throttled_time"]
            - record["precpu_stats"]["throttling_data"]["throttled_time"]
        ) / 1000000

        j.append(tick)
    return j


def plot(df, title="foo", output=None):
    fig = make_subplots(
        rows=2, cols=1, shared_xaxes=True, subplot_titles=("CPU %", "CPU (trottled) ms")
    )
    fig.add_trace(
        go.Scatter(x=df["read"], y=df["cpu_percent"], name="CPU %"),
        row=1,
        col=1,
    )
    fig.add_trace(
        go.Scatter(x=df["read"], y=df["cpu_delta"], name="CPU ms"),
        row=2,
        col=1,
    )
    fig.add_trace(
        go.Scatter(x=df["read"], y=df["cpu_throttled_ms"], name="throttled ms"),
        row=2,
        col=1,
    )

    fig.update_xaxes(title_text="Time")
    fig.update_yaxes(title_text="%", range=[0, 100], row=1, col=1)
    fig.update_yaxes(title_text="ms", row=2, col=1)
    fig.update_layout(title_text=title)

    if output is None:
        fig.show()
    else:
        fig.write_html(output)


if __name__ == "__main__":
    """
    Generate the input file with something like:
    curl -s --unix-socket /var/run/docker.sock \
        "http://localhost/containers/${CONTAINER_ID}/stats?stream=1" |
            jq -c  '{read: .read, cpu_stats: .cpu_stats, precpu_stats: .precpu_stats}'
    """
    parser = argparse.ArgumentParser(description="Plot docker stats CPU usage data")
    parser.add_argument("input", type=argparse.FileType("r"))
    parser.add_argument("--out", type=argparse.FileType("w"), required=False)
    args = parser.parse_args()

    j = read_json(args.input)
    df = pd.DataFrame.from_dict(j)
    plot(df, args.input.name, args.out)