#!/usr/bin/env python3

# Copyright 2020 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.

# Library to extract scenario definitions from scenario_config.py.
#
# Contains functions to filter, analyze and dump scenario definitions.
#
# This library is used in loadtest_config.py to generate the "scenariosJSON"
# field in the format accepted by the OSS benchmarks framework.
# See https://github.com/grpc/test-infra/blob/master/config/samples/cxx_example_loadtest.yaml
#
# It can also be used to dump scenarios to files, to count scenarios by
# language, and to export scenario languages in a format that can be used for
# automation.
#
# Example usage:
#
# scenario_config.py --export_scenarios -l cxx -f cxx_scenario_ -r '.*' \
#     --category=scalable
#
# scenario_config.py --count_scenarios
#
# scenario_config.py --count_scenarios --category=scalable
#
# For usage of the language config output, see loadtest_config.py.

import argparse
import collections
import json
import re
import sys
from typing import Any, Callable, Dict, Iterable, NamedTuple

import scenario_config

# Language parameters for load test config generation.

LanguageConfig = NamedTuple(
    "LanguageConfig",
    [
        ("category", str),
        ("language", str),
        ("client_language", str),
        ("server_language", str),
    ],
)


def category_string(categories: Iterable[str], category: str) -> str:
    """Converts a list of categories into a single string for counting."""
    if category != "all":
        return category if category in categories else ""

    main_categories = ("scalable", "smoketest")
    s = set(categories)

    c = [m for m in main_categories if m in s]
    s.difference_update(main_categories)
    c.extend(s)
    return " ".join(c)


def gen_scenario_languages(category: str) -> Iterable[LanguageConfig]:
    """Generates tuples containing the languages specified in each scenario."""
    for language in scenario_config.LANGUAGES:
        for scenario in scenario_config.LANGUAGES[language].scenarios():
            client_language = scenario.get("CLIENT_LANGUAGE", "")
            server_language = scenario.get("SERVER_LANGUAGE", "")
            categories = scenario.get("CATEGORIES", [])
            if category != "all" and category not in categories:
                continue
            cat = category_string(categories, category)
            yield LanguageConfig(
                category=cat,
                language=language,
                client_language=client_language,
                server_language=server_language,
            )


def scenario_filter(
    scenario_name_regex: str = ".*",
    category: str = "all",
    client_language: str = "",
    server_language: str = "",
) -> Callable[[Dict[str, Any]], bool]:
    """Returns a function to filter scenarios to process."""

    def filter_scenario(scenario: Dict[str, Any]) -> bool:
        """Filters scenarios that match specified criteria."""
        if not re.search(scenario_name_regex, scenario["name"]):
            return False
        # if the 'CATEGORIES' key is missing, treat scenario as part of
        # 'scalable' and 'smoketest'. This matches the behavior of
        # run_performance_tests.py.
        scenario_categories = scenario.get(
            "CATEGORIES", ["scalable", "smoketest"]
        )
        if category not in scenario_categories and category != "all":
            return False

        scenario_client_language = scenario.get("CLIENT_LANGUAGE", "")
        if client_language != scenario_client_language:
            return False

        scenario_server_language = scenario.get("SERVER_LANGUAGE", "")
        if server_language != scenario_server_language:
            return False

        return True

    return filter_scenario


def gen_scenarios(
    language_name: str,
    scenario_filter_function: Callable[[Dict[str, Any]], bool],
) -> Iterable[Dict[str, Any]]:
    """Generates scenarios that match a given filter function."""
    return map(
        scenario_config.remove_nonproto_fields,
        filter(
            scenario_filter_function,
            scenario_config.LANGUAGES[language_name].scenarios(),
        ),
    )


def dump_to_json_files(
    scenarios: Iterable[Dict[str, Any]], filename_prefix: str
) -> None:
    """Dumps a list of scenarios to JSON files"""
    count = 0
    for scenario in scenarios:
        filename = "{}{}.json".format(filename_prefix, scenario["name"])
        print("Writing file {}".format(filename), file=sys.stderr)
        with open(filename, "w") as outfile:
            # The dump file should have {"scenarios" : []} as the top level
            # element, when embedded in a LoadTest configuration YAML file.
            json.dump({"scenarios": [scenario]}, outfile, indent=2)
        count += 1
    print("Wrote {} scenarios".format(count), file=sys.stderr)


def main() -> None:
    language_choices = sorted(scenario_config.LANGUAGES.keys())
    argp = argparse.ArgumentParser(description="Exports scenarios to files.")
    argp.add_argument(
        "--export_scenarios",
        action="store_true",
        help="Export scenarios to JSON files.",
    )
    argp.add_argument(
        "--count_scenarios",
        action="store_true",
        help="Count scenarios for all test languages.",
    )
    argp.add_argument(
        "-l", "--language", choices=language_choices, help="Language to export."
    )
    argp.add_argument(
        "-f",
        "--filename_prefix",
        default="scenario_dump_",
        type=str,
        help="Prefix for exported JSON file names.",
    )
    argp.add_argument(
        "-r",
        "--regex",
        default=".*",
        type=str,
        help="Regex to select scenarios to run.",
    )
    argp.add_argument(
        "--category",
        default="all",
        choices=[
            "all",
            "inproc",
            "scalable",
            "smoketest",
            "sweep",
            "psm",
            "dashboard",
        ],
        help="Select scenarios for a category of tests.",
    )
    argp.add_argument(
        "--client_language",
        default="",
        choices=language_choices,
        help="Select only scenarios with a specified client language.",
    )
    argp.add_argument(
        "--server_language",
        default="",
        choices=language_choices,
        help="Select only scenarios with a specified server language.",
    )
    args = argp.parse_args()

    if args.export_scenarios and not args.language:
        print(
            "Dumping scenarios requires a specified language.", file=sys.stderr
        )
        argp.print_usage(file=sys.stderr)
        return

    if args.export_scenarios:
        s_filter = scenario_filter(
            scenario_name_regex=args.regex,
            category=args.category,
            client_language=args.client_language,
            server_language=args.server_language,
        )
        scenarios = gen_scenarios(args.language, s_filter)
        dump_to_json_files(scenarios, args.filename_prefix)

    if args.count_scenarios:
        print(
            "Scenario count for all languages (category: {}):".format(
                args.category
            )
        )
        print(
            "{:>5}  {:16} {:8} {:8} {}".format(
                "Count", "Language", "Client", "Server", "Categories"
            )
        )
        c = collections.Counter(gen_scenario_languages(args.category))
        total = 0
        for (cat, l, cl, sl), count in c.most_common():
            print(
                "{count:5}  {l:16} {cl:8} {sl:8} {cat}".format(
                    l=l, cl=cl, sl=sl, count=count, cat=cat
                )
            )
            total += count

        print(
            "\n{:>5}  total scenarios (category: {})".format(
                total, args.category
            )
        )


if __name__ == "__main__":
    main()