# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

"""
General utility functions for AWS.
This is for very simple, general functions. Use facades for service-specific macro operations.
These function may or may not require an AWS account in order to use.
"""

import itertools
import logging
import time
from urllib.parse import urlparse

from botocore.exceptions import ClientError

from czc import constants

logger = logging.getLogger(constants.PROJECT_NAME)

PERMISSION_EXCEPTIONS = [
    'UnauthorizedOperation',
    'AccessDenied',
    'AuthorizationError',
    'AccessDeniedException'
]

THROTTLE_EXCEPTIONS = [
    'ThrottlingException',
    'Throttling',
    'TooManyRequestsException',
    'RequestLimitExceeded',
    'ThrottledException'
]

PAGINATION_TOKEN_KEYS = ['nextToken', 'Marker', 'NextToken', 'NextMarker']
PAGINATION_THROTTLE_BACKOFF_TIME = 20


def _get_paginator(api, client, kwargs, last_page=None):
    logger.debug('Getting paginator')
    if client.can_paginate(api):
        if last_page:
            next_token = _next_pagination_token(last_page)
            if 'PaginationConfig' in kwargs:
                if not isinstance(kwargs['PaginationConfig'], dict):
                    kwargs['PaginationConfig'] = {}
                kwargs['PaginationConfig']['StartingToken'] = next_token
            else:
                kwargs['PaginationConfig'] = {'StartingToken': next_token}

        return client.get_paginator(api).paginate(**kwargs)

    return _get_fallback_paginator(client, api, last_page, kwargs)


def _done_paginating(page):
    for key in PAGINATION_TOKEN_KEYS:
        if key in page:
            return False

    return True


def _token_key(page):
    for key in PAGINATION_TOKEN_KEYS:
        if key in page:
            return key

    # NOTE: a current limitation of the generic fallback paginator
    # is that it can't differentiate between an API whose result payload is missing a
    # token because there is only one page of results, and an API whose result is
    # a token because it doesn't support pagination.
    # in both of these cases, the fallback will only return the first page of results.
    return None


def _next_pagination_token(page):
    key = _token_key(page)
    return page.get(key)


def get_paginated_results(client, api: str, item_key: str, **kwargs: dict) -> list:
    """
    The AWS API limits the number of items it includes in a response
    to a maximum of 50. For API Calls that produce more results than can
    be returned at once, AWS includes a ContinuationToken in its response
    which can be supplied to subsequent API calls to tell Amazon where you
    left off.
    This method employs boto3's paginators (which are an abstraction of the
    ContinuationToken concept described above) to call AWS repeatedly,
    collecting the outputs until there are no more results. It returns
    all the results it collected to the caller.
    NOTE: This applies to boto3 service clients only; services which
    have boto3 resources have built-in support for paginators.
    Args:
        client: (boto3 client) - A client instance for which to use a paginator
        api: (str) - The name of an API for which boto3 defines a paginator,
             corresponding to the currently instantiated service
        item_key: (str) - The name of the key containing the list metadata in the client's response, e.g. 'Items'
        kwargs: Key Word Arguments to pass through to API call
    """

    def pages():
        paginator_iterator = _get_paginator(api, client, kwargs)

        page = None
        is_done_paginating = False
        while not is_done_paginating:
            logger.debug('Getting another page')
            try:
                for page in paginator_iterator:
                    is_done_paginating = _done_paginating(page)
                    yield page
            except ClientError as err:
                error_code = err.response['Error']['Code']
                logger.debug(f'Error code during pagination {error_code}')
                if error_code in THROTTLE_EXCEPTIONS:
                    logger.warning(f'Pagination for api {api} throttled, sleeping 20 seconds')
                    time.sleep(PAGINATION_THROTTLE_BACKOFF_TIME)
                    paginator_iterator = _get_paginator(api, client, kwargs, last_page=page)
                else:
                    raise err

    client_responses = [
        p.get(item_key, [])
        for p in pages()
    ]
    return list(itertools.chain(*client_responses))


def _get_fallback_paginator(client, api, start_page, kwargs):
    # if boto doesn't have a custom paginator,
    # attempt a more generic pagination strategy.
    # there are a few strategies that are pervasive throughout amazon.
    try:
        invokable_api = getattr(client, api)
    except AttributeError:
        msg = f'the api ({api}) for which pagination was requested does not exist '
        logger.exception(msg)
        raise

    if not start_page:
        start_page = invokable_api(**kwargs) if invokable_api else {}
        yield start_page

    pagination_token_field = _token_key(start_page)
    pagination_token = start_page.get(pagination_token_field)

    while pagination_token:
        page = invokable_api({**kwargs, pagination_token_field: pagination_token})
        pagination_token = page.get(pagination_token_field)
        yield page


def parse_s3_url(s3_url):
    """
    Parse an S3 URL, returning a tuple containing the bucket name and key

    Args:
        s3_url (string): s3 URL, like s3://some_bucket/some_path/and/maybe/a.file

    Returns (tuple): bucket name (some_bucket) and key (some_path/and/maybe/a.file)

    >>> parse_s3_url("s3://some_bucket/some_path/and/maybe/a.file")
    ('some_bucket', 'some_path/and/maybe/a.file')

    >>> parse_s3_url("s3://some_bucket/")
    ('some_bucket', '')

    >>> parse_s3_url("sdfsdfsdfsf")
    Traceback (most recent call last):
    ...
    aws.MalformedS3Url: S3 url 'sdfsdfsdfsf' is malformed

    """
    parsed_config_location = urlparse(s3_url)
    if parsed_config_location.scheme == 's3':
        bucket = parsed_config_location.netloc
        key = parsed_config_location.path.lstrip('/')
    else:
        raise MalformedS3Url(f"S3 url '{s3_url}' is malformed")

    return (bucket, key)


class MalformedS3Url(Exception):
    """Thrown when we have a malformed S3 URL"""
    pass
