#!/usr/bin/env python3
# Copyright 2017 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Utilities for manipulating JSON data that represents microbenchmark results.

import os

# template arguments and dynamic arguments of individual benchmark types
# Example benchmark name: "BM_UnaryPingPong<TCP, NoOpMutator, NoOpMutator>/0/0"
_BM_SPECS = {
    "BM_UnaryPingPong": {
        "tpl": ["fixture", "client_mutator", "server_mutator"],
        "dyn": ["request_size", "response_size"],
    },
    "BM_PumpStreamClientToServer": {
        "tpl": ["fixture"],
        "dyn": ["request_size"],
    },
    "BM_PumpStreamServerToClient": {
        "tpl": ["fixture"],
        "dyn": ["request_size"],
    },
    "BM_StreamingPingPong": {
        "tpl": ["fixture", "client_mutator", "server_mutator"],
        "dyn": ["request_size", "request_count"],
    },
    "BM_StreamingPingPongMsgs": {
        "tpl": ["fixture", "client_mutator", "server_mutator"],
        "dyn": ["request_size"],
    },
    "BM_PumpStreamServerToClient_Trickle": {
        "tpl": [],
        "dyn": ["request_size", "bandwidth_kilobits"],
    },
    "BM_PumpUnbalancedUnary_Trickle": {
        "tpl": [],
        "dyn": ["cli_req_size", "svr_req_size", "bandwidth_kilobits"],
    },
    "BM_ErrorStringOnNewError": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_ErrorStringRepeatedly": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_ErrorGetStatus": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_ErrorGetStatusCode": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_ErrorHttpError": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_HasClearGrpcStatus": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_IsolatedFilter": {
        "tpl": ["fixture", "client_mutator"],
        "dyn": [],
    },
    "BM_HpackEncoderEncodeHeader": {
        "tpl": ["fixture"],
        "dyn": ["end_of_stream", "request_size"],
    },
    "BM_HpackParserParseHeader": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_CallCreateDestroy": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_Zalloc": {
        "tpl": [],
        "dyn": ["request_size"],
    },
    "BM_PollEmptyPollset_SpeedOfLight": {
        "tpl": [],
        "dyn": ["request_size", "request_count"],
    },
    "BM_StreamCreateSendInitialMetadataDestroy": {
        "tpl": ["fixture"],
        "dyn": [],
    },
    "BM_TransportStreamSend": {
        "tpl": [],
        "dyn": ["request_size"],
    },
    "BM_TransportStreamRecv": {
        "tpl": [],
        "dyn": ["request_size"],
    },
    "BM_StreamingPingPongWithCoalescingApi": {
        "tpl": ["fixture", "client_mutator", "server_mutator"],
        "dyn": ["request_size", "request_count", "end_of_stream"],
    },
    "BM_Base16SomeStuff": {
        "tpl": [],
        "dyn": ["request_size"],
    },
}


def numericalize(s):
    """Convert abbreviations like '100M' or '10k' to a number."""
    if not s:
        return ""
    if s[-1] == "k":
        return float(s[:-1]) * 1024
    if s[-1] == "M":
        return float(s[:-1]) * 1024 * 1024
    if 0 <= (ord(s[-1]) - ord("0")) <= 9:
        return float(s)
    assert "not a number: %s" % s


def parse_name(name):
    cpp_name = name
    if "<" not in name and "/" not in name and name not in _BM_SPECS:
        return {"name": name, "cpp_name": name}
    rest = name
    out = {}
    tpl_args = []
    dyn_args = []
    if "<" in rest:
        tpl_bit = rest[rest.find("<") + 1 : rest.rfind(">")]
        arg = ""
        nesting = 0
        for c in tpl_bit:
            if c == "<":
                nesting += 1
                arg += c
            elif c == ">":
                nesting -= 1
                arg += c
            elif c == ",":
                if nesting == 0:
                    tpl_args.append(arg.strip())
                    arg = ""
                else:
                    arg += c
            else:
                arg += c
        tpl_args.append(arg.strip())
        rest = rest[: rest.find("<")] + rest[rest.rfind(">") + 1 :]
    if "/" in rest:
        s = rest.split("/")
        rest = s[0]
        dyn_args = s[1:]
    name = rest
    assert name in _BM_SPECS, "_BM_SPECS needs to be expanded for %s" % name
    assert len(dyn_args) == len(_BM_SPECS[name]["dyn"])
    assert len(tpl_args) == len(_BM_SPECS[name]["tpl"])
    out["name"] = name
    out["cpp_name"] = cpp_name
    out.update(
        dict(
            (k, numericalize(v))
            for k, v in zip(_BM_SPECS[name]["dyn"], dyn_args)
        )
    )
    out.update(dict(zip(_BM_SPECS[name]["tpl"], tpl_args)))
    return out


def expand_json(js):
    if not js:
        raise StopIteration()
    for bm in js["benchmarks"]:
        if bm["name"].endswith("_stddev") or bm["name"].endswith("_mean"):
            continue
        context = js["context"]
        if "label" in bm:
            labels_list = [
                s.split(":")
                for s in bm["label"].strip().split(" ")
                if len(s) and s[0] != "#"
            ]
            for el in labels_list:
                el[0] = el[0].replace("/iter", "_per_iteration")
            labels = dict(labels_list)
        else:
            labels = {}
        # TODO(jtattermusch): grabbing kokoro env values shouldn't be buried
        # deep in the JSON conversion logic.
        # Link the data to a kokoro job run by adding
        # well known kokoro env variables as metadata for each row
        row = {
            "jenkins_build": os.environ.get("KOKORO_BUILD_NUMBER", ""),
            "jenkins_job": os.environ.get("KOKORO_JOB_NAME", ""),
        }
        row.update(context)
        row.update(bm)
        row.update(parse_name(row["name"]))
        row.update(labels)
        yield row