#!/usr/bin/env python3
# Copyright 2021 The 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.

# This script generates a load test configuration template from a collection of
# load test configurations.
#
# Configuration templates contain client and server configurations for multiple
# languages, and may contain template substitution keys. These templates are
# used to generate load test configurations by selecting clients and servers for
# the required languages. The source files for template generation may be load
# test configurations or load test configuration templates. Load test
# configuration generation is performed by loadtest_config.py. See documentation
# below:
# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md

import argparse
import os
import sys
from typing import Any, Dict, Iterable, List, Mapping, Type

import yaml

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import loadtest_config

TEMPLATE_FILE_HEADER_COMMENT = """
# Template generated from load test configurations by loadtest_template.py.
#
# Configuration templates contain client and server configurations for multiple
# languages, and may contain template substitution keys. These templates are
# used to generate load test configurations by selecting clients and servers for
# the required languages. The source files for template generation may be load
# test configurations or load test configuration templates. Load test
# configuration generation is performed by loadtest_config.py. See documentation
# below:
# https://github.com/grpc/grpc/blob/master/tools/run_tests/performance/README.md
"""


def insert_worker(
    worker: Dict[str, Any], workers: List[Dict[str, Any]]
) -> None:
    """Inserts client or server into a list, without inserting duplicates."""

    def dump(w):
        return yaml.dump(w, Dumper=yaml.SafeDumper, default_flow_style=False)

    worker_str = dump(worker)
    if any((worker_str == dump(w) for w in workers)):
        return
    workers.append(worker)


def uniquify_workers(workermap: Dict[str, List[Dict[str, Any]]]) -> None:
    """Name workers if there is more than one for the same map key."""
    for workers in list(workermap.values()):
        if len(workers) <= 1:
            continue
        for i, worker in enumerate(workers):
            worker["name"] = str(i)


def loadtest_template(
        input_file_names: Iterable[str],
        metadata: Mapping[str, Any],
        inject_client_pool: bool,
        inject_driver_image: bool,
        inject_driver_pool: bool,
        inject_server_pool: bool,
        inject_big_query_table: bool,
        inject_timeout_seconds: bool,
        inject_ttl_seconds: bool) -> Dict[str, Any]:  # fmt: skip
    """Generates the load test template."""
    spec = dict()  # type: Dict[str, Any]
    clientmap = dict()  # Dict[str, List[Dict[str, Any]]]
    servermap = dict()  # Dict[Str, List[Dict[str, Any]]]
    template = {
        "apiVersion": "e2etest.grpc.io/v1",
        "kind": "LoadTest",
        "metadata": metadata,
    }
    for input_file_name in input_file_names:
        with open(input_file_name) as f:
            input_config = yaml.safe_load(f.read())

            if input_config.get("apiVersion") != template["apiVersion"]:
                raise ValueError(
                    "Unexpected api version in file {}: {}".format(
                        input_file_name, input_config.get("apiVersion")
                    )
                )
            if input_config.get("kind") != template["kind"]:
                raise ValueError(
                    "Unexpected kind in file {}: {}".format(
                        input_file_name, input_config.get("kind")
                    )
                )

            for client in input_config["spec"]["clients"]:
                del client["name"]
                if inject_client_pool:
                    client["pool"] = "${client_pool}"
                if client["language"] not in clientmap:
                    clientmap[client["language"]] = []
                insert_worker(client, clientmap[client["language"]])

            for server in input_config["spec"]["servers"]:
                del server["name"]
                if inject_server_pool:
                    server["pool"] = "${server_pool}"
                if server["language"] not in servermap:
                    servermap[server["language"]] = []
                insert_worker(server, servermap[server["language"]])

            input_spec = input_config["spec"]
            del input_spec["clients"]
            del input_spec["servers"]
            del input_spec["scenariosJSON"]
            spec.update(input_config["spec"])

    uniquify_workers(clientmap)
    uniquify_workers(servermap)

    spec.update(
        {
            "clients": sum(
                (clientmap[language] for language in sorted(clientmap)),
                start=[],
            ),
            "servers": sum(
                (servermap[language] for language in sorted(servermap)),
                start=[],
            ),
        }
    )

    if "driver" not in spec:
        spec["driver"] = {"language": "cxx"}

    driver = spec["driver"]
    if "name" in driver:
        del driver["name"]
    if inject_driver_image:
        if "run" not in driver:
            driver["run"] = [{"name": "main"}]
        driver["run"][0]["image"] = "${driver_image}"
    if inject_driver_pool:
        driver["pool"] = "${driver_pool}"

    if "run" not in driver:
        if inject_driver_pool:
            raise ValueError("Cannot inject driver.pool: missing driver.run.")
        del spec["driver"]

    if inject_big_query_table:
        if "results" not in spec:
            spec["results"] = dict()
        spec["results"]["bigQueryTable"] = "${big_query_table}"
    if inject_timeout_seconds:
        spec["timeoutSeconds"] = "${timeout_seconds}"
    if inject_ttl_seconds:
        spec["ttlSeconds"] = "${ttl_seconds}"

    template["spec"] = spec

    return template


def template_dumper(header_comment: str) -> Type[yaml.SafeDumper]:
    """Returns a custom dumper to dump templates in the expected format."""

    class TemplateDumper(yaml.SafeDumper):
        def expect_stream_start(self):
            super().expect_stream_start()
            if isinstance(self.event, yaml.StreamStartEvent):
                self.write_indent()
                self.write_indicator(header_comment, need_whitespace=False)

    def str_presenter(dumper, data):
        if "\n" in data:
            return dumper.represent_scalar(
                "tag:yaml.org,2002:str", data, style="|"
            )
        return dumper.represent_scalar("tag:yaml.org,2002:str", data)

    TemplateDumper.add_representer(str, str_presenter)

    return TemplateDumper


def main() -> None:
    argp = argparse.ArgumentParser(
        description="Creates a load test config generator template.",
        fromfile_prefix_chars="@",
    )
    argp.add_argument(
        "-i",
        "--inputs",
        action="extend",
        nargs="+",
        type=str,
        help="Input files.",
    )
    argp.add_argument(
        "-o",
        "--output",
        type=str,
        help="Output file. Outputs to stdout if not set.",
    )
    argp.add_argument(
        "--inject_client_pool",
        action="store_true",
        help="Set spec.client(s).pool values to '${client_pool}'.",
    )
    argp.add_argument(
        "--inject_driver_image",
        action="store_true",
        help="Set spec.driver(s).image values to '${driver_image}'.",
    )
    argp.add_argument(
        "--inject_driver_pool",
        action="store_true",
        help="Set spec.driver(s).pool values to '${driver_pool}'.",
    )
    argp.add_argument(
        "--inject_server_pool",
        action="store_true",
        help="Set spec.server(s).pool values to '${server_pool}'.",
    )
    argp.add_argument(
        "--inject_big_query_table",
        action="store_true",
        help="Set spec.results.bigQueryTable to '${big_query_table}'.",
    )
    argp.add_argument(
        "--inject_timeout_seconds",
        action="store_true",
        help="Set spec.timeoutSeconds to '${timeout_seconds}'.",
    )
    argp.add_argument(
        "--inject_ttl_seconds", action="store_true", help="Set timeout "
    )
    argp.add_argument(
        "-n", "--name", default="", type=str, help="metadata.name."
    )
    argp.add_argument(
        "-a",
        "--annotation",
        action="append",
        type=str,
        help="metadata.annotation(s), in the form key=value.",
        dest="annotations",
    )
    args = argp.parse_args()

    annotations = loadtest_config.parse_key_value_args(args.annotations)

    metadata = {"name": args.name}
    if annotations:
        metadata["annotations"] = annotations

    template = loadtest_template(
        input_file_names=args.inputs,
        metadata=metadata,
        inject_client_pool=args.inject_client_pool,
        inject_driver_image=args.inject_driver_image,
        inject_driver_pool=args.inject_driver_pool,
        inject_server_pool=args.inject_server_pool,
        inject_big_query_table=args.inject_big_query_table,
        inject_timeout_seconds=args.inject_timeout_seconds,
        inject_ttl_seconds=args.inject_ttl_seconds,
    )

    with open(args.output, "w") if args.output else sys.stdout as f:
        yaml.dump(
            template,
            stream=f,
            Dumper=template_dumper(TEMPLATE_FILE_HEADER_COMMENT.strip()),
            default_flow_style=False,
        )


if __name__ == "__main__":
    main()