# Copyright 2019 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.
"""Decorators for use with caching library """
import json


class InjectSecretString:
    """Decorator implementing high-level Secrets Manager caching client"""

    def __init__(self, secret_id, cache):
        """
        Constructs a decorator to inject a single non-keyworded argument from a cached secret for a given function.

        :type secret_id: str
        :param secret_id: The secret identifier

        :type cache: aws_secretsmanager_caching.SecretCache
        :param cache: Secret cache
        """

        self.cache = cache
        self.secret_id = secret_id

    def __call__(self, func):
        """
        Return a function with cached secret injected as first argument.

        :type func: object
        :param func: The function for injecting a single non-keyworded argument too.
        :return The function with the injected argument.
        """

        secret = self.cache.get_secret_string(secret_id=self.secret_id)

        def _wrapped_func(*args, **kwargs):
            """
            Internal function to execute wrapped function
            """
            return func(secret, *args, **kwargs)

        return _wrapped_func


class InjectKeywordedSecretString:
    """Decorator implementing high-level Secrets Manager caching client using JSON-based secrets"""

    def __init__(self, secret_id, cache, **kwargs):
        """
        Construct a decorator to inject a variable list of keyword arguments to a given function with resolved values
        from a cached secret.

        :type kwargs: dict
        :param kwargs: dictionary mapping original keyword argument of wrapped function to JSON-encoded secret key

        :type secret_id: str
        :param secret_id: The secret identifier

        :type cache: aws_secretsmanager_caching.SecretCache
        :param cache: Secret cache
        """

        self.cache = cache
        self.kwarg_map = kwargs
        self.secret_id = secret_id

    def __call__(self, func):
        """
        Return a function with injected keyword arguments from a cached secret.

        :type func: object
        :param func: function for injecting keyword arguments.
        :return The original function with injected keyword arguments
        """

        try:
            secret = json.loads(self.cache.get_secret_string(secret_id=self.secret_id))
        except json.decoder.JSONDecodeError:
            raise RuntimeError('Cached secret is not valid JSON') from None

        resolved_kwargs = dict()
        for orig_kwarg in self.kwarg_map:
            secret_key = self.kwarg_map[orig_kwarg]
            try:
                resolved_kwargs[orig_kwarg] = secret[secret_key]
            except KeyError:
                raise RuntimeError('Cached secret does not contain key {0}'.format(secret_key)) from None

        def _wrapped_func(*args, **kwargs):
            """
            Internal function to execute wrapped function
            """
            return func(*args, **resolved_kwargs, **kwargs)

        return _wrapped_func