# Copyright 2013 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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 has customizations to unify paging parameters.

For any operation that can be paginated, we will:

    * Hide the service specific pagination params.  This can vary across
    services and we're going to replace them with a consistent set of
    arguments.  The arguments will still work, but they are not
    documented.  This allows us to add a pagination config after
    the fact and still remain backwards compatible with users that
    were manually doing pagination.
    * Add a ``--starting-token`` and a ``--max-items`` argument.

"""

import logging
import sys
from functools import partial

from botocore import model, xform_name
from botocore.exceptions import DataNotFoundError

from awscli.arguments import BaseCLIArgument
from awscli.customizations.exceptions import ParamValidationError
from awscli.customizations.utils import uni_print

logger = logging.getLogger(__name__)


STARTING_TOKEN_HELP = """
<p>A token to specify where to start paginating.  This is the
<code>NextToken</code> from a previously truncated response.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""

MAX_ITEMS_HELP = """
<p>The total number of items to return in the command's output.
If the total number of items available is more than the value
specified, a <code>NextToken</code> is provided in the command's
output.  To resume pagination, provide the
<code>NextToken</code> value in the <code>starting-token</code>
argument of a subsequent command.  <b>Do not</b> use the
<code>NextToken</code> response element directly outside of the
AWS CLI.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""

PAGE_SIZE_HELP = """
<p>The size of each page to get in the AWS service call.  This
does not affect the number of items returned in the command's
output.  Setting a smaller page size results in more calls to
the AWS service, retrieving fewer items in each call.  This can
help prevent the AWS service calls from timing out.</p>
<p>For usage examples, see <a
href="https://docs.aws.amazon.com/cli/latest/userguide/pagination.html"
>Pagination</a> in the <i>AWS Command Line Interface User
Guide</i>.</p>
"""


def register_pagination(event_handlers):
    event_handlers.register('building-argument-table', unify_paging_params)
    event_handlers.register_last('doc-description', add_paging_description)


def get_paginator_config(session, service_name, operation_name):
    try:
        paginator_model = session.get_paginator_model(service_name)
    except DataNotFoundError:
        return None
    try:
        operation_paginator_config = paginator_model.get_paginator(
            operation_name
        )
    except ValueError:
        return None
    return operation_paginator_config


def add_paging_description(help_command, **kwargs):
    # This customization is only applied to the description of
    # Operations, so we must filter out all other events.
    if not isinstance(help_command.obj, model.OperationModel):
        return
    service_name = help_command.obj.service_model.service_name
    paginator_config = get_paginator_config(
        help_command.session, service_name, help_command.obj.name
    )
    if not paginator_config:
        return
    help_command.doc.style.new_paragraph()
    help_command.doc.writeln(
        (
            '``%s`` is a paginated operation. Multiple API calls may be issued '
            'in order to retrieve the entire data set of results. You can '
            'disable pagination by providing the ``--no-paginate`` argument.'
        )
        % help_command.name
    )
    # Only include result key information if it is present.
    if paginator_config.get('result_key'):
        queries = paginator_config['result_key']
        if type(queries) is not list:
            queries = [queries]
        queries = ", ".join([('``%s``' % s) for s in queries])
        help_command.doc.writeln(
            (
                'When using ``--output text`` and the ``--query`` argument on a '
                'paginated response, the ``--query`` argument must extract data '
                'from the results of the following query expressions: %s'
            )
            % queries
        )


def unify_paging_params(
    argument_table, operation_model, event_name, session, **kwargs
):
    paginator_config = get_paginator_config(
        session,
        operation_model.service_model.service_name,
        operation_model.name,
    )
    if paginator_config is None:
        # We only apply these customizations to paginated responses.
        return
    logger.debug(
        "Modifying paging parameters for operation: %s", operation_model.name
    )
    _remove_existing_paging_arguments(argument_table, paginator_config)
    parsed_args_event = event_name.replace(
        'building-argument-table.', 'operation-args-parsed.'
    )
    call_parameters_event = event_name.replace(
        'building-argument-table', 'calling-command'
    )
    shadowed_args = {}
    add_paging_argument(
        argument_table,
        'starting-token',
        PageArgument(
            'starting-token',
            STARTING_TOKEN_HELP,
            parse_type='string',
            serialized_name='StartingToken',
        ),
        shadowed_args,
    )
    input_members = operation_model.input_shape.members
    type_name = 'integer'
    if 'limit_key' in paginator_config:
        limit_key_shape = input_members[paginator_config['limit_key']]
        type_name = limit_key_shape.type_name
        if type_name not in PageArgument.type_map:
            raise TypeError(
                (
                    'Unsupported pagination type {0} for operation {1}'
                    ' and parameter {2}'
                ).format(
                    type_name,
                    operation_model.name,
                    paginator_config['limit_key'],
                )
            )
        add_paging_argument(
            argument_table,
            'page-size',
            PageArgument(
                'page-size',
                PAGE_SIZE_HELP,
                parse_type=type_name,
                serialized_name='PageSize',
            ),
            shadowed_args,
        )

    add_paging_argument(
        argument_table,
        'max-items',
        PageArgument(
            'max-items',
            MAX_ITEMS_HELP,
            parse_type=type_name,
            serialized_name='MaxItems',
        ),
        shadowed_args,
    )
    # We will register two pagination handlers.
    #
    # The first is focused on analyzing the CLI arguments passed to see
    # if they contain explicit operation-specific pagination args.  If so,
    # they are doing manual pagination and we turn off the CLI pagination.
    #
    # The second is called later in the event chain and analyzes the actual
    # calling parameters passed to the operation.
    #
    # The reason we have to do the second is that someone could put
    # operation-specific pagination args in the CLI input JSON file
    # directly and this bypasses all of the CLI args processing.
    session.register(
        parsed_args_event,
        partial(
            check_should_enable_pagination,
            list(_get_all_cli_input_tokens(paginator_config)),
            shadowed_args,
            argument_table,
        ),
    )
    session.register(
        call_parameters_event,
        partial(
            check_should_enable_pagination_call_parameters,
            list(_get_all_input_tokens(paginator_config)),
        ),
    )


def add_paging_argument(argument_table, arg_name, argument, shadowed_args):
    if arg_name in argument_table:
        # If there's already an entry in the arg table for this argument,
        # this means we're shadowing an argument for this operation.  We
        # need to store this later in case pagination is turned off because
        # we put these arguments back.
        # See the comment in check_should_enable_pagination() for more info.
        shadowed_args[arg_name] = argument_table[arg_name]
    argument_table[arg_name] = argument


def check_should_enable_pagination(
    input_tokens,
    shadowed_args,
    argument_table,
    parsed_args,
    parsed_globals,
    **kwargs,
):
    normalized_paging_args = ['start_token', 'max_items']
    for token in input_tokens:
        py_name = token.replace('-', '_')
        if (
            getattr(parsed_args, py_name) is not None
            and py_name not in normalized_paging_args
        ):
            # The user has specified a manual (undocumented) pagination arg.
            # We need to automatically turn pagination off.
            logger.debug(
                "User has specified a manual pagination arg. "
                "Automatically setting --no-paginate."
            )
            parsed_globals.paginate = False

    if not parsed_globals.paginate:
        ensure_paging_params_not_set(parsed_args, shadowed_args)
        # Because pagination is now disabled, there's a chance that
        # we were shadowing arguments.  For example, we inject a
        # --max-items argument in unify_paging_params().  If the
        # the operation also provides its own MaxItems (which we
        # expose as --max-items) then our custom pagination arg
        # was shadowing the customers arg.  When we turn pagination
        # off we need to put back the original argument which is
        # what we're doing here.
        for key, value in shadowed_args.items():
            argument_table[key] = value


def ensure_paging_params_not_set(parsed_args, shadowed_args):
    paging_params = ['starting_token', 'page_size', 'max_items']
    shadowed_params = [p.replace('-', '_') for p in shadowed_args.keys()]
    params_used = [
        p
        for p in paging_params
        if p not in shadowed_params and getattr(parsed_args, p, None)
    ]

    if len(params_used) > 0:
        converted_params = ', '.join(
            ["--" + p.replace('_', '-') for p in params_used]
        )
        raise ParamValidationError(
            "Cannot specify --no-paginate along with pagination "
            "arguments: %s" % converted_params
        )


def _remove_existing_paging_arguments(argument_table, pagination_config):
    for cli_name in _get_all_cli_input_tokens(pagination_config):
        argument_table[cli_name]._UNDOCUMENTED = True


def _get_all_cli_input_tokens(pagination_config):
    # Get all input tokens including the limit_key
    # if it exists.
    tokens = _get_input_tokens(pagination_config)
    for token_name in tokens:
        cli_name = xform_name(token_name, '-')
        yield cli_name
    if 'limit_key' in pagination_config:
        key_name = pagination_config['limit_key']
        cli_name = xform_name(key_name, '-')
        yield cli_name


# Get all tokens but return them in API namespace rather than CLI namespace
def _get_all_input_tokens(pagination_config):
    # Get all input tokens including the limit_key
    # if it exists.
    tokens = _get_input_tokens(pagination_config)
    for token_name in tokens:
        yield token_name
    if 'limit_key' in pagination_config:
        key_name = pagination_config['limit_key']
        yield key_name


def _get_input_tokens(pagination_config):
    tokens = pagination_config['input_token']
    if not isinstance(tokens, list):
        return [tokens]
    return tokens


def _get_cli_name(param_objects, token_name):
    for param in param_objects:
        if param.name == token_name:
            return param.cli_name.lstrip('-')


# This function checks for pagination args in the actual calling
# arguments passed to the function.  If the user is using the
# --cli-input-json parameter to provide JSON parameters they are
# all in the API naming space rather than the CLI naming space
# and would be missed by the processing above.  This function gets
# called on the calling-command event.
def check_should_enable_pagination_call_parameters(
    input_tokens, call_parameters, parsed_args, parsed_globals, **kwargs
):
    for param in call_parameters:
        if param in input_tokens:
            logger.debug(
                "User has specified a manual pagination arg. "
                "Automatically setting --no-paginate."
            )
            parsed_globals.paginate = False


class PageArgument(BaseCLIArgument):
    type_map = {
        'string': str,
        'integer': int,
        'long': int,
    }

    def __init__(self, name, documentation, parse_type, serialized_name):
        self.argument_model = model.Shape('PageArgument', {'type': 'string'})
        self._name = name
        self._serialized_name = serialized_name
        self._documentation = documentation
        self._parse_type = parse_type
        self._required = False

    def _emit_non_positive_max_items_warning(self):
        uni_print(
            "warning: Non-positive values for --max-items may result in undefined behavior.\n",
            sys.stderr,
        )

    @property
    def cli_name(self):
        return '--' + self._name

    @property
    def cli_type_name(self):
        return self._parse_type

    @property
    def required(self):
        return self._required

    @required.setter
    def required(self, value):
        self._required = value

    @property
    def documentation(self):
        return self._documentation

    def add_to_parser(self, parser):
        parser.add_argument(
            self.cli_name,
            dest=self.py_name,
            type=self.type_map[self._parse_type],
        )

    def add_to_params(self, parameters, value):
        if value is not None:
            if self._serialized_name == 'MaxItems' and int(value) <= 0:
                self._emit_non_positive_max_items_warning()
            pagination_config = parameters.get('PaginationConfig', {})
            pagination_config[self._serialized_name] = value
            parameters['PaginationConfig'] = pagination_config