import os.path
import glob
import json


KW_NODES = "nodes"
KW_ID = "id"  # uniq, use in edges
KW_LABEL = "label"  # short name
KW_DETAIL = "detail"  # show when hover
KW_TYPE = "type"

KW_EDGES = "edges"
KW_SRC = "source"
KW_DST = "target"


def from_loghub(path):
    files = glob.glob(os.path.join(path, "sched*"))
    jsons = []
    for p in files:
        with open(p) as f:
            jsons.append(json.load(f))

    return trans(jsons)


def fmt_duraion(s):
    i = int(s)
    r = ""
    for a, b in [(86400, 'd'), (3600, 'h'), (60, 'm')]:
        if i >= a:
            r += ("{}" + b).format(i / a)
            i %= a
    if r == "":
        # return "{:.1e}s".format(s)
        return "{}s".format(i + 1)
    else:
        return r


def M(b):
    return int(b/(1024*1024))


def summary_prof(p):
    counters = p['counters']
    stats = p['stats']
    info = p['info']

    task = counters['task']
    fail = counters['fail']
    mem = stats['bytes_max_rss']
    t = stats['secs_all']
    task_torun = task['all'] - task['running'] - task['finished']
    fail_error = fail['all'] - fail['oom'] - fail['timeout']
    finished = task['finished']
    if info['finish_time'] > 0:
        time_stage = info['finish_time'] - info['start_time']
    else:
        time_stage = "None"
    if finished > 0:
        avg_mem = mem['sum'] / finished
        avg_time = t['sum'] / finished
        speedup = avg_time * finished / time_stage
    else:
        avg_mem = 0
        avg_time = 0
        speedup = 0

    return [
        ["task", "{} = {} + {} + {}".format(task['all'], task['finished'], task['running'], task_torun)],
        ["fail", "{} = {} + {} + {}".format(fail['all'], fail['oom'],  fail_error, fail['timeout'])],
        ['mem',  "{} || [{}, {}, {}]".format(info['mem'], M(mem['min']), M(avg_mem), M(mem['max']))],
        ['time', "{} || [{}, {}, {}]".format(*[fmt_duraion(s) for s in [time_stage, t['min'], avg_time, t['max']]])],
        ['speedup',  "{:.2f}".format(speedup)]
    ]


def trans(runs):
    api_nodes = {}
    api_edges = {}
    stage_nodes = {}
    stage_edges = {}
    for r in runs:
        r = r["run"]
        for s in r['stages']:
            for n in s['graph'][KW_NODES]:
                if n[KW_ID] in stage_nodes:
                    continue
                rdds = n['rdds']
                n['rdds'] = list(reversed([{"k": rdd["rdd_name"], "v": str(rdd["api_callsite_id"])}
                                           for rdd in rdds]))
                if n[KW_ID] == s['info']['output_pipeline']:
                    p = n['prof'] = {
                        'info': s['info'],
                        'stats': s['stats'],
                        'counters': s['counters']
                    }
                    n['prof_summary'] = summary_prof(p)
                    n['is_output'] = True
                else:
                    n['is_output'] = False

                stage_nodes[n[KW_ID]] = n
            for e in s['graph'][KW_EDGES]:
                id_ = e[KW_SRC], e[KW_DST]
                if id_ not in stage_edges:
                    e['info'] = ",".join(["{}={}".format(str(k), str(v)) for k, v in e["info"].items()])
                    stage_edges[id_] = e
        sink_node = r["sink"]['node']
        sink_node['call_id'] = str(sink_node['call_id'])
        stage_nodes[sink_node[KW_ID]] = sink_node
        sink_edge = r['sink'][KW_EDGES]
        stage_edges[(sink_edge[KW_SRC], sink_edge[KW_DST])] = sink_edge

        c = r['call_graph']
        for n in c[KW_NODES]:
            id_ = n[KW_ID] = n['call_id'] = str(n[KW_ID])
            api_nodes[id_] = n
        for e in c[KW_EDGES]:
            s, r = str(e[KW_SRC]), str(e[KW_DST])
            if (s, r) not in api_edges:
                e[KW_SRC] = s
                e[KW_DST] = r
                api_edges[(s, r)] = e
    res = {
        "stages": {
            KW_NODES: list(stage_nodes.values()),
            KW_EDGES: list(stage_edges.values()),
        },
        "calls": {
            KW_NODES: list(api_nodes.values()),
            KW_EDGES: list(api_edges.values()),
        }
    }

    return res

@@ -935,67 +935,9 @@
 l'%5D%0A
-    mem = stats%5B'bytes_max_rss'%5D%0A    t = stats%5B'secs_all'%5D
 %0A
+
     
@@ -1087,24 +1087,348 @@
 'finished'%5D%0A
+    res = %5B%0A        %5B%22task%22, %22%7B%7D = %7B%7D + %7B%7D + %7B%7D%22.format(task%5B'all'%5D, task%5B'finished'%5D, task%5B'running'%5D, task_torun)%5D,%0A        %5B%22fail%22, %22%7B%7D = %7B%7D + %7B%7D + %7B%7D%22.format(fail%5B'all'%5D, fail%5B'oom'%5D, fail_error, fail%5B'timeout'%5D)%5D,%0A    %5D%0A%0A    if not stats:%0A        return res%0A%0A    mem = stats%5B'bytes_max_rss'%5D%0A    t = stats%5B'secs_all'%5D%0A%0A
     if info%5B
@@ -1780,222 +1780,14 @@
   re
-turn %5B%0A        %5B%22task%22, %22%7B%7D = %7B%7D + %7B%7D + %7B%7D%22.format(task%5B'all'%5D, task%5B'finished'%5D, task%5B'running'%5D, task_torun)%5D,%0A        %5B%22fail%22, %22%7B%7D = %7B%7D + %7B%7D + %7B%7D%22.format(fail%5B'all'%5D, fail%5B'oom'%5D,  fail_error, fail%5B'timeout'%5D)%5D,
+s2 = %5B
 %0A   
@@ -2050,22 +2050,44 @@
 eedup)%5D%0A
-
     %5D%0A
+    return res + res2%0A
 %0A%0Adef tr

deal with empty task stats.