#!/usr/bin/env python3
################################################################################
# Copyright 2020-2024 Intel Corporation
#
# 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.
################################################################################

import sys

import argparse
from argparse import RawTextHelpFormatter

from src import utils
from src import writer


def convert(
    verbose_level,
    parser,
    input,
    action,
    generator,
    split_output,
    agg_keys,
    events=["create", "exec"],
):
    status = utils.check_version()
    if status != utils.status.get("SUCCESS"):
        return status

    logger = writer.Writer(verbose_level=verbose_level)
    log_parser = None
    if parser == "oneDNN":
        from src import dnnl_parser

        log_parser = dnnl_parser.LogParser(logger, input)
    else:
        logger.print("Error: unsupported parser", "STDIO")
        return utils.status.get("FAILED")

    logger.print(f"Processing input ...", "INFO")
    log_parser.process(events)

    output = None
    if action == "dumpIR":
        logger.print(f"Dumping data from input...", "INFO")
        log_parser.dump(True)

    if action == "generate":
        logger.print(f"Generating output ...", "INFO")
        if generator == "benchdnn":
            from src import benchdnn_generator

            gen = benchdnn_generator.InputGenerator(logger)
            output = gen.generate(log_parser.get_data(), split_output)
        elif generator == "breakdown":
            from src import breakdown_generator

            gen = breakdown_generator.BreakdownGenerator(logger)
            output = gen.generate(log_parser.get_data(), agg_keys)
        else:
            logger.print("Error: unsupported generator", "STDIO")
            return utils.status.get("FAILED")

    return utils.status.get("SUCCESS"), output


def validate_option(value, supported_values, str):
    if not value in supported_values:
        print(f"ERROR: {str}")
        return utils.status.get("FAILED")
    return utils.status.get("SUCCESS")


def main():
    status = utils.check_version()
    if status != utils.status.get("SUCCESS"):
        return status

    action_opts = ["generate", "dumpIR"]
    generator_opts = ["benchdnn", "breakdown"]
    parser_opts = ["oneDNN"]
    verbose_opts = ["0", "1"]
    aggregate_opts = [
        "engine",
        "prim_kind",
        "impl",
        "prop_kind",
        "mds",
        "exts",
        "aux",
        "shapes",
    ]
    event_opts = ["exec", "create"]
    args_parser = argparse.ArgumentParser(
        description="oneDNN log converter", formatter_class=RawTextHelpFormatter
    )
    args_parser.add_argument(
        "-i", "--input", default="stdin", help="input file (default: stdin)"
    )
    args_parser.add_argument(
        "-p",
        "--parser",
        default="oneDNN",
        help=f"type of parser (default: oneDNN). Values: {parser_opts}.",
    )
    args_parser.add_argument(
        "-a",
        "--action",
        default="generate",
        help=f"an action (default: generate). Values: {action_opts}.",
    )
    args_parser.add_argument(
        "-s",
        "--split",
        type=bool,
        default=False,
        help="split generated inputs by primitive kinds (default: False)",
    )
    args_parser.add_argument(
        "-k",
        "--aggregate",
        nargs="+",
        default=aggregate_opts,
        help=f"aggregates statistics on the specified keys (default: all keys but time).\nValues: {aggregate_opts}",
    )
    args_parser.add_argument(
        "-v",
        "--verbose_level",
        default="0",
        help=f"verbose level (default: 0). Values: {verbose_opts}.",
    )
    args_parser.add_argument(
        "-o", "--output", default="stdout", help="output file (default: stdout)"
    )
    args_parser.add_argument(
        "-g",
        "--generator",
        default="benchdnn",
        help=f"target generator (default: benchdnn). Values: {generator_opts}.",
    )
    args_parser.add_argument(
        "-e",
        "--events",
        nargs="+",
        default=event_opts,
        help=f"events to parse (default: create and exec).\nValues: {event_opts}.",
    )
    args = args_parser.parse_args()

    # validate options
    status = validate_option(args.action, action_opts, "Unknown action value")
    if status != utils.status.get("SUCCESS"):
        return status
    status = validate_option(
        args.verbose_level, verbose_opts, "Unknown verbose_level value"
    )
    if status != utils.status.get("SUCCESS"):
        return status
    status = validate_option(args.parser, parser_opts, "Unknown parser value")
    if status != utils.status.get("SUCCESS"):
        return status
    status = validate_option(args.generator, generator_opts, "Unknown generator value")
    if status != utils.status.get("SUCCESS"):
        return status

    input_data = []
    if args.input == "stdin":
        # if no input was piped, skip reading
        if not sys.stdin.isatty():
            for line in sys.stdin:
                input_data.append(line)
        else:
            print("WARN: no input was provided to the script")
            args_parser.print_help()
    else:
        try:
            input_data = open(args.input, "r").readlines()
        except BaseException as e:
            print(f"Error while reading input: {e}")

    output = None

    event_sets = args.events if args.generator == 'breakdown' else [args.events]

    for events in event_sets:
        status, output = convert(
            verbose_level=args.verbose_level,
            parser=args.parser,
            input=input_data,
            action=args.action,
            generator=args.generator,
            split_output=args.split,
            agg_keys=args.aggregate,
            events=events
        )

        if status != utils.status.get("SUCCESS"):
            return status

        if output != None:
            if args.output != "stdout":
                if output != None:
                    for key, value in output.items():
                        filename = args.output
                        if args.split == True:
                            filename += "." + key
                        of = open(filename, "w")
                    if args.generator == "breakdown":
                        print(f"Event: {events}", file=of)
                    print(value, end="", file=of)
            else:
                if args.generator == "breakdown":
                    print(f"Event: {events}")
                for key, value in output.items():
                    if args.split == False:
                        print(f"{value}")
                    else:
                        print(f"--{key}\n{value}")


if __name__ == "__main__":
    main()
