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 name_from_field(field: str) -> str :
    if field.startswith("onlyoffice-"):
        return "onlyoffice"
    return field.split("_")[0]

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 data[0]}
    print("Found", len(data), "points")


    for t_i, stat in enumerate(data):
        for field in stat:
            if field == "date":  # date
                data_dict[field][t_i] = datetime.strptime(stat[field], "%Y-%m-%dT%H:%M:%S%z")
            else:  # float
                value = stat[field].split(" ")[0]
                value = value.replace("MiB", "e3")
                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):
    bar_chart = pygal.StackedBar(height=400)
    bar_chart.x_labels = data["date"]
    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()