import argparse
import json
from datetime import datetime

import pygal


PNG_OUTPUT = 'mem_usage.png'


def parse_args():
    parser = argparse.ArgumentParser(
        description=("Render for docker stats memory usage")
    )

    parser.add_argument('stats', type=str, help="Path to stats file")
    parser.add_argument('--web', '-w', action='store_true',
                        help="Render in web browser instead of svg")
    args = parser.parse_args()
    return args.stats, args.web


def extract_names(data: list):
    names = set()
    for d in data:
        for field in d:
            name = name_from_field(field)
            names.update([name])
    return names


def name_from_field(field: str) -> str :
    if field.startswith("onlyoffice-"):
        return "onlyoffice"
    name = field.split("_")[0]

    name = name.replace("org-caracals-", "")
    name = name.replace(".caracals.org", "")
    return name


def main():
    stats_fn, web_render = parse_args()
    data = load_data(stats_fn)

    render(data, web_render)


def load_data(stats_fn: str):
    with open(stats_fn) as stats_f:
        data = json.load(stats_f)

    data_dict = {name_from_field(field): [0]*len(data) for field in extract_names(data)}

    print("Found", len(data), "points")
    print("  keys:", data_dict.keys())

    for t_i, stat in enumerate(data):
        for field in stat:
            if field == "date":  # date
                date_ = stat[field].replace("+01:00", "+0000").replace("+02:00", "+0200")
                data_dict[field][t_i] = datetime.strptime(date_, "%Y-%m-%dT%H:%M:%S%z")
            else:  # float
                value = stat[field].split(" ")[0]
                value = value.replace("MiB", "e3")
                value = value.replace("GiB", "e6")
                value = value.replace("B", "")
                data_dict[name_from_field(field)][t_i] += float(value) / 1000 # values are in MiB
    return data_dict


def render(data: dict, web_render:bool=False):
    style = pygal.style.Style(value_label_font_size=5, value_font_size=5, label_font_size=5, legend_font_size=5)
    bar_chart = pygal.StackedBar(height=400, x_label_rotation=25, style=style, legend_box_size=5)
    labels = [d.strftime("%B %d %H:%M") for d in data["date"]]
    bar_chart.x_labels = labels
    for k in data:
        if k == "date":
            continue
        #if "db" in k or "database" in k or "mysql" in k or "mongo" in k or "postgre" in k:
        #    continue
        bar_chart.add(k, data[k])

    if web_render:
        bar_chart.render_in_browser()
    else:
        bar_chart.render_to_png(PNG_OUTPUT)
        print("Image generated in:", PNG_OUTPUT)

if __name__ == "__main__":
    main()