# Copyright 2018 The Abseil 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 module provides argparse integration with absl.flags.

argparse_flags.ArgumentParser is a drop-in replacement for
argparse.ArgumentParser. It takes care of collecting and defining absl flags
in argparse.

Here is a simple example:

    # Assume the following absl.flags is defined in another module:
    #
    #     from absl import flags
    #     flags.DEFINE_string('echo', None, 'The echo message.')
    #
    parser = argparse_flags.ArgumentParser(
        description='A demo of absl.flags and argparse integration.')
    parser.add_argument('--header', help='Header message to print.')

    # The parser will also accept the absl flag `--echo`.
    # The `header` value is available as `args.header` just like a regular
    # argparse flag. The absl flag `--echo` continues to be available via
    # `absl.flags.FLAGS` if you want to access it.
    args = parser.parse_args()

    # Example usages:
    # ./program --echo='A message.' --header='A header'
    # ./program --header 'A header' --echo 'A message.'


Here is another example demonstrates subparsers:

    parser = argparse_flags.ArgumentParser(description='A subcommands demo.')
    parser.add_argument('--header', help='The header message to print.')

    subparsers = parser.add_subparsers(help='The command to execute.')

    roll_dice_parser = subparsers.add_parser(
        'roll_dice', help='Roll a dice.',
        # By default, absl flags can also be specified after the sub-command.
        # To only allow them before sub-command, pass
        # `inherited_absl_flags=None`.
        inherited_absl_flags=None)
    roll_dice_parser.add_argument('--num_faces', type=int, default=6)
    roll_dice_parser.set_defaults(command=roll_dice)

    shuffle_parser = subparsers.add_parser('shuffle', help='Shuffle inputs.')
    shuffle_parser.add_argument(
        'inputs', metavar='I', nargs='+', help='Inputs to shuffle.')
    shuffle_parser.set_defaults(command=shuffle)

    args = parser.parse_args(argv[1:])
    args.command(args)

    # Example usages:
    # ./program --echo='A message.' roll_dice --num_faces=6
    # ./program shuffle --echo='A message.' 1 2 3 4


There are several differences between absl.flags and argparse_flags:

1. Flags defined with absl.flags are parsed differently when using the
   argparse parser. Notably:

   1) absl.flags allows both single-dash and double-dash for any flag, and
      doesn't distinguish them; argparse_flags only allows double-dash for
      flag's regular name, and single-dash for flag's `short_name`.
   2) Boolean flags in absl.flags can be specified with `--bool`, `--nobool`,
      as well as `--bool=true/false` (though not recommended);
      in argparse_flags, it only allows `--bool`, `--nobool`.

2. Help related flag differences:
   1) absl.flags does not define help flags, absl.app does that; argparse_flags
      defines help flags unless passed with `add_help=False`.
   2) absl.app supports `--helpxml`; argparse_flags does not.
   3) argparse_flags supports `-h`; absl.app does not.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import sys

from absl import flags


_BUILT_IN_FLAGS = frozenset({
    'help',
    'helpshort',
    'helpfull',
    'helpxml',
    'flagfile',
    'undefok',
})


class ArgumentParser(argparse.ArgumentParser):
  """Custom ArgumentParser class to support special absl flags."""

  def __init__(self, **kwargs):
    """Initializes ArgumentParser.

    Args:
      **kwargs: same as argparse.ArgumentParser, except:
          1. It also accepts `inherited_absl_flags`: the absl flags to inherit.
             The default is the global absl.flags.FLAGS instance. Pass None to
             ignore absl flags.
          2. The `prefix_chars` argument must be the default value '-'.

    Raises:
      ValueError: Raised when prefix_chars is not '-'.
    """
    prefix_chars = kwargs.get('prefix_chars', '-')
    if prefix_chars != '-':
      raise ValueError(
          'argparse_flags.ArgumentParser only supports "-" as the prefix '
          'character, found "{}".'.format(prefix_chars))

    # Remove inherited_absl_flags before calling super.
    self._inherited_absl_flags = kwargs.pop('inherited_absl_flags', flags.FLAGS)
    # Now call super to initialize argparse.ArgumentParser before calling
    # add_argument in _define_absl_flags.
    super(ArgumentParser, self).__init__(**kwargs)

    if self.add_help:
      # -h and --help are defined in super.
      # Also add the --helpshort and --helpfull flags.
      self.add_argument(
          # Action 'help' defines a similar flag to -h/--help.
          '--helpshort', action='help',
          default=argparse.SUPPRESS, help=argparse.SUPPRESS)
      self.add_argument(
          '--helpfull', action=_HelpFullAction,
          default=argparse.SUPPRESS, help='show full help message and exit')

    if self._inherited_absl_flags:
      self.add_argument('--undefok', help=argparse.SUPPRESS)
      self._define_absl_flags(self._inherited_absl_flags)

  def parse_known_args(self, args=None, namespace=None):
    if args is None:
      args = sys.argv[1:]
    if self._inherited_absl_flags:
      # Handle --flagfile.
      # Explicitly specify force_gnu=True, since argparse behaves like
      # gnu_getopt: flags can be specified after positional arguments.
      args = self._inherited_absl_flags.read_flags_from_files(
          args, force_gnu=True)

    undefok_missing = object()
    undefok = getattr(namespace, 'undefok', undefok_missing)

    namespace, args = super(ArgumentParser, self).parse_known_args(
        args, namespace)

    # For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where
    # sub-parsers don't preserve existing namespace attributes.
    # Restore the undefok attribute if a sub-parser dropped it.
    if undefok is not undefok_missing:
      namespace.undefok = undefok

    if self._inherited_absl_flags:
      # Handle --undefok. At this point, `args` only contains unknown flags,
      # so it won't strip defined flags that are also specified with --undefok.
      # For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where
      # sub-parsers don't preserve existing namespace attributes. The undefok
      # attribute might not exist because a subparser dropped it.
      if hasattr(namespace, 'undefok'):
        args = _strip_undefok_args(namespace.undefok, args)
        # absl flags are not exposed in the Namespace object. See Namespace:
        # https://docs.python.org/3/library/argparse.html#argparse.Namespace.
        del namespace.undefok
      self._inherited_absl_flags.mark_as_parsed()
      try:
        self._inherited_absl_flags.validate_all_flags()
      except flags.IllegalFlagValueError as e:
        self.error(str(e))

    return namespace, args

  def _define_absl_flags(self, absl_flags):
    """Defines flags from absl_flags."""
    key_flags = set(absl_flags.get_key_flags_for_module(sys.argv[0]))
    for name in absl_flags:
      if name in _BUILT_IN_FLAGS:
        # Do not inherit built-in flags.
        continue
      flag_instance = absl_flags[name]
      # Each flags with short_name appears in FLAGS twice, so only define
      # when the dictionary key is equal to the regular name.
      if name == flag_instance.name:
        # Suppress the flag in the help short message if it's not a main
        # module's key flag.
        suppress = flag_instance not in key_flags
        self._define_absl_flag(flag_instance, suppress)

  def _define_absl_flag(self, flag_instance, suppress):
    """Defines a flag from the flag_instance."""
    flag_name = flag_instance.name
    short_name = flag_instance.short_name
    argument_names = ['--' + flag_name]
    if short_name:
      argument_names.insert(0, '-' + short_name)
    if suppress:
      helptext = argparse.SUPPRESS
    else:
      # argparse help string uses %-formatting. Escape the literal %'s.
      helptext = flag_instance.help.replace('%', '%%')
    if flag_instance.boolean:
      # Only add the `no` form to the long name.
      argument_names.append('--no' + flag_name)
      self.add_argument(
          *argument_names, action=_BooleanFlagAction, help=helptext,
          metavar=flag_instance.name.upper(),
          flag_instance=flag_instance)
    else:
      self.add_argument(
          *argument_names, action=_FlagAction, help=helptext,
          metavar=flag_instance.name.upper(),
          flag_instance=flag_instance)


class _FlagAction(argparse.Action):
  """Action class for Abseil non-boolean flags."""

  def __init__(self, option_strings, dest, help, metavar, flag_instance):  # pylint: disable=redefined-builtin
    """Initializes _FlagAction.

    Args:
      option_strings: See argparse.Action.
      dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS.
      help: See argparse.Action.
      metavar: See argparse.Action.
      flag_instance: absl.flags.Flag, the absl flag instance.
    """
    del dest
    self._flag_instance = flag_instance
    super(_FlagAction, self).__init__(
        option_strings=option_strings,
        dest=argparse.SUPPRESS,
        help=help,
        metavar=metavar)

  def __call__(self, parser, namespace, values, option_string=None):
    """See https://docs.python.org/3/library/argparse.html#action-classes."""
    self._flag_instance.parse(values)
    self._flag_instance.using_default_value = False


class _BooleanFlagAction(argparse.Action):
  """Action class for Abseil boolean flags."""

  def __init__(self, option_strings, dest, help, metavar, flag_instance):  # pylint: disable=redefined-builtin
    """Initializes _BooleanFlagAction.

    Args:
      option_strings: See argparse.Action.
      dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS.
      help: See argparse.Action.
      metavar: See argparse.Action.
      flag_instance: absl.flags.Flag, the absl flag instance.
    """
    del dest
    self._flag_instance = flag_instance
    flag_names = [self._flag_instance.name]
    if self._flag_instance.short_name:
      flag_names.append(self._flag_instance.short_name)
    self._flag_names = frozenset(flag_names)
    super(_BooleanFlagAction, self).__init__(
        option_strings=option_strings,
        dest=argparse.SUPPRESS,
        nargs=0,  # Does not accept values, only `--bool` or `--nobool`.
        help=help,
        metavar=metavar)

  def __call__(self, parser, namespace, values, option_string=None):
    """See https://docs.python.org/3/library/argparse.html#action-classes."""
    if not isinstance(values, list) or values:
      raise ValueError('values must be an empty list.')
    if option_string.startswith('--'):
      option = option_string[2:]
    else:
      option = option_string[1:]
    if option in self._flag_names:
      self._flag_instance.parse('true')
    else:
      if not option.startswith('no') or option[2:] not in self._flag_names:
        raise ValueError('invalid option_string: ' + option_string)
      self._flag_instance.parse('false')
    self._flag_instance.using_default_value = False


class _HelpFullAction(argparse.Action):
  """Action class for --helpfull flag."""

  def __init__(self, option_strings, dest, default, help):  # pylint: disable=redefined-builtin
    """Initializes _HelpFullAction.

    Args:
      option_strings: See argparse.Action.
      dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS.
      default: Ignored.
      help: See argparse.Action.
    """
    del dest, default
    super(_HelpFullAction, self).__init__(
        option_strings=option_strings,
        dest=argparse.SUPPRESS,
        default=argparse.SUPPRESS,
        nargs=0,
        help=help)

  def __call__(self, parser, namespace, values, option_string=None):
    """See https://docs.python.org/3/library/argparse.html#action-classes."""
    # This only prints flags when help is not argparse.SUPPRESS.
    # It includes user defined argparse flags, as well as main module's
    # key absl flags. Other absl flags use argparse.SUPPRESS, so they aren't
    # printed here.
    parser.print_help()

    absl_flags = parser._inherited_absl_flags  # pylint: disable=protected-access
    if absl_flags:
      modules = sorted(absl_flags.flags_by_module_dict())
      main_module = sys.argv[0]
      if main_module in modules:
        # The main module flags are already printed in parser.print_help().
        modules.remove(main_module)
      print(absl_flags._get_help_for_modules(  # pylint: disable=protected-access
          modules, prefix='', include_special_flags=True))
    parser.exit()


def _strip_undefok_args(undefok, args):
  """Returns a new list of args after removing flags in --undefok."""
  if undefok:
    undefok_names = set(name.strip() for name in undefok.split(','))
    undefok_names |= set('no' + name for name in undefok_names)
    # Remove undefok flags.
    args = [arg for arg in args if not _is_undefok(arg, undefok_names)]
  return args


def _is_undefok(arg, undefok_names):
  """Returns whether we can ignore arg based on a set of undefok flag names."""
  if not arg.startswith('-'):
    return False
  if arg.startswith('--'):
    arg_without_dash = arg[2:]
  else:
    arg_without_dash = arg[1:]
  if '=' in arg_without_dash:
    name, _ = arg_without_dash.split('=', 1)
  else:
    name = arg_without_dash
  if name in undefok_names:
    return True
  return False