#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 argparse
import sys
from textwrap import dedent

import polygraphy
from polygraphy import mod
from polygraphy.logger import G_LOGGER
from polygraphy.tools.args import LoggerArgs
from polygraphy.tools.args.base import ArgGroups


@mod.export()
class Tool:
    """
    Base class for CLI Tools.
    """

    def __init__(self, name=None):
        self.name = name
        self.arg_groups = (
            ArgGroups()
        )  # Populated by setup_parser based on get_subscriptions()

    def setup_parser(self, subparsers=None):
        """
        Set up a command-line argument parser.

        Args:
            subparsers (argparse.SubParsers):
                    A subparser group from argparse, like that returned by ``ArgumentParser.add_subparsers()``.
                    If this is omitted, this function will generate a new ``ArgumentParser`` instance.
                    Defaults to None.

        Returns:
            argparse.ArgumentParser:
                    The newly created parser if ``subparsers`` is not provided, or the newly created subparser otherwise.
        """
        assert self.__doc__, "No help output was provided for this tool!"

        subscriptions = self.get_subscriptions()
        # Always subscribe to the logger arguments first.
        for arg_group in [LoggerArgs()] + subscriptions:
            m_type = type(arg_group)
            self.arg_groups[m_type] = arg_group

        allow_abbrev = all(
            arg_group.allows_abbreviation() for arg_group in self.arg_groups.values()
        )

        description = dedent(self.__doc__)
        if subparsers is not None:
            summary = []
            for line in description.strip().splitlines():
                if line.isspace() or not line:
                    break
                summary.append(line)
            summary = "\n".join(summary)

            parser = subparsers.add_parser(
                self.name,
                help=summary,
                add_help=True,
                description=description,
                allow_abbrev=allow_abbrev,
                formatter_class=argparse.RawDescriptionHelpFormatter,
            )
            parser.set_defaults(subcommand=self)
        else:
            parser = argparse.ArgumentParser(
                add_help=True, description=description, allow_abbrev=allow_abbrev
            )

        for arg_group in self.arg_groups.values():
            arg_group.register(self.arg_groups)
            # This must be done after registration, since some argument groups
            # may conditionally define arguments based on what other groups are present.
            arg_group.add_parser_args(parser)

        try:
            self.add_parser_args(parser)
        except Exception as err:
            G_LOGGER.internal_error(
                f"Could not register tool argument parser for: {self.name}\nNote: Error was: {err}"
            )
        return parser

    # Implementation for `get_subscriptions`. This should be implemented by child classes instead of `get_subscriptions`
    def get_subscriptions_impl(self):
        return []

    def get_subscriptions(self):
        """
        Returns the list of argument groups this tools wishes to subscribe to.

        Returns:
            List[BaseArgs]: The list of argument groups to subscribe to.
        """
        return self.get_subscriptions_impl()

    # Implementation for `get_subtools`. This should be implemented by child classes instead of `get_subtools`
    def get_subtools_impl(self):
        return "", []

    def get_subtools(self):
        """
        Returns the subtools of this tool.

        Returns:
            Tuple[str, List[Tool]]:
                    A tuple containing the title to use for the subtools
                    and a list of subtools.
        """
        return self.get_subtools_impl()

    # Implementation for `add_parser_args`. This should be implemented by child classes instead of `add_parser_args`
    def add_parser_args_impl(self, parser):
        pass

    def add_parser_args(self, parser):
        """
        Add arguments to the command-line parser.
        This will also add arguments for any subtools reported by `get_subtools()`.

        Args:
            parser (argparse.ArgumentParser): The argument parser.
        """
        self.add_parser_args_impl(parser)

        title, subtools = self.get_subtools()
        if subtools:
            subparsers = parser.add_subparsers(title=title, dest="subtool")
            subparsers.required = True

            for subtool in subtools:
                subtool.setup_parser(subparsers)

    def parse(self, args):
        """
        Run argument parsing for all argument groups.

        Args:
            args (Namespace):
                    The namespace returned by ``parse_args()`` or ``parse_known_args()``.
        """
        for arg_group in self.arg_groups.values():
            arg_group.parse(args)

    # Implementation for `run`. This should be implemented by child classes instead of `run`
    def run_impl(self, args):
        raise NotImplementedError("run_impl() must be implemented by child classes")

    def run(self, args):
        """
        Runs this tool with the specified arguments.
        This must be called only after calling `parse()`.

        Args:
            args (Namespace):
                    The namespace returned by ``parse_args()`` or ``parse_known_args()``.
        """
        G_LOGGER.module_info(polygraphy)

        status = self.run_impl(args)
        if status is None:
            status = 0
        return status

    # Implementation for `show_start_end_logging`. This should be implemented by child classes instead of `show_start_end_logging`
    def show_start_end_logging_impl(self, args):
        return False

    def show_start_end_logging(self, args):
        """
        Whether to display start and end times before and after running the tools.

        Args:
            args (Namespace):
                    The namespace returned by ``parse_args()`` or ``parse_known_args()``.

        Returns:
            bool: Whether to display start/end times.
        """
        return self.show_start_end_logging_impl(args)

    def main(self):
        """
        Set up and run this tool. This function serves as a replacement for a manually
        defined ``main`` method.

        Runs ``sys.exit()`` with the status code returned by ``run``. If ``run`` does
        not return anything, always exits with ``0`` (success).
        """
        parser = self.setup_parser()
        args = parser.parse_args()
        self.parse(args)
        sys.exit(self.run(args))
