"""Module containing extension app authenticator classes.
"""
from abc import ABC, abstractmethod

from symphony.bdk.core.auth.auth_session import AppAuthSession
from symphony.bdk.core.auth.jwt_helper import validate_jwt, create_signed_jwt
from symphony.bdk.core.auth.tokens_repository import TokensRepository, InMemoryTokensRepository
from symphony.bdk.core.config.model.bdk_retry_config import BdkRetryConfig
from symphony.bdk.core.config.model.bdk_rsa_key_config import BdkRsaKeyConfig
from symphony.bdk.core.retry import retry
from symphony.bdk.core.retry.strategy import authentication_retry
from symphony.bdk.gen.auth_api.certificate_authentication_api import CertificateAuthenticationApi
from symphony.bdk.gen.auth_api.certificate_pod_api import CertificatePodApi
from symphony.bdk.gen.auth_model.extension_app_authenticate_request import ExtensionAppAuthenticateRequest
from symphony.bdk.gen.login_api.authentication_api import AuthenticationApi
from symphony.bdk.gen.login_model.authenticate_extension_app_request import AuthenticateExtensionAppRequest
from symphony.bdk.gen.login_model.extension_app_tokens import ExtensionAppTokens
from symphony.bdk.gen.pod_api.pod_api import PodApi
from symphony.bdk.gen.pod_model.pod_certificate import PodCertificate


class ExtensionAppAuthenticator(ABC):
    """Base abstract class to handle extension app authentication.
    """

    def __init__(self, app_id: str, tokens_repository: TokensRepository = None):
        """

        :param app_id: the application ID
        :param tokens_repository: the tokens repository to store existing valid sessions.
          Defaults to InMemoryTokensRepository
        """
        self._app_id = app_id
        self._tokens_repository = tokens_repository or InMemoryTokensRepository()

    async def authenticate_extension_app(self, app_token: str) -> AppAuthSession:
        """Authenticates an extension app.

        :param app_token: the token generated by the application backend
        :return: the extension app authentication session
        """
        auth_session = AppAuthSession(self, app_token)
        await auth_session.refresh()

        return auth_session

    async def validate_jwt(self, jwt: str) -> dict:
        """Validates a jwt against the pod certificate.

        :param jwt: the jwt to be validated
        :return: the dictionary of jwt claims
        :raise AuthInitializationError: If the pod certificate or jwt are invalid.
        """
        pod_certificate = await self._get_pod_certificate()
        return validate_jwt(jwt, pod_certificate.certificate, self._app_id)

    async def authenticate_and_retrieve_tokens(self, app_token: str) -> ExtensionAppTokens:
        """Actually authenticates and retrieves the tokens.
        Save the tokens in the tokens repository.

        :param app_token: the token generated by the application backend
        :return: the extension app tokens
        """
        ext_app_tokens = await self._retrieve_tokens(app_token)
        await self._tokens_repository.save(ext_app_tokens.app_token, ext_app_tokens.symphony_token)

        return ext_app_tokens

    async def is_token_pair_valid(self, app_token: str, symphony_token: str) -> bool:
        """Validates if appToken and symphonyToken corresponds to an existing session.

        :param app_token: the application token
        :param symphony_token: the Symphony token
        :return: True if token pair is valid, False otherwise
        """
        retrieved_sym_token = await self._tokens_repository.get(app_token)
        return retrieved_sym_token is not None and retrieved_sym_token == symphony_token

    @abstractmethod
    async def _get_pod_certificate(self) -> PodCertificate:
        """Retrieves the pod certificate. Used when validating JWT tokens.

        :return: the pod certificate
        """

    @abstractmethod
    async def _retrieve_tokens(self, app_token: str) -> ExtensionAppTokens:
        """Actually authenticates and retrieves the tokens.

        :param app_token: the token generated by the application backend
        :return: the extension app tokens
        """


class ExtensionAppAuthenticatorRsa(ExtensionAppAuthenticator):
    """A subclass of :class:`ExtensionAppAuthenticator` specific to extension app RSA authentication.
    """

    def __init__(self,
                 authentication_api: AuthenticationApi,
                 pod_api: PodApi,
                 app_id: str,
                 private_key_config: BdkRsaKeyConfig,
                 retry_config: BdkRetryConfig,
                 tokens_repository: TokensRepository = None):
        """

        :param authentication_api: the AuthenticationApi instance
        :param pod_api: the PodApi instance
        :param app_id: the application ID
        :param private_key_config: the private key configuration of the extension app
        :param retry_config: retry configuration
        :param tokens_repository: the tokens repository to store existing valid sessions.
          Defaults to InMemoryTokensRepository
        """
        super().__init__(app_id, tokens_repository)
        self._authentication_api = authentication_api
        self._pod_api = pod_api
        self._private_key_config = private_key_config
        self._retry_config = retry_config

    @retry(retry=authentication_retry)
    async def _retrieve_tokens(self, app_token: str) -> ExtensionAppTokens:
        jwt = create_signed_jwt(self._private_key_config, self._app_id)
        authentication_request = AuthenticateExtensionAppRequest(app_token=app_token, auth_token=jwt)

        return await self._authentication_api.v1_pubkey_app_authenticate_extension_app_post(authentication_request)

    @retry(retry=authentication_retry)
    async def _get_pod_certificate(self) -> PodCertificate:
        return await self._pod_api.v1_podcert_get()


class ExtensionAppAuthenticatorCert(ExtensionAppAuthenticator):
    """A subclass of :class:`ExtensionAppAuthenticator` specific to extension app certificate authentication.
    """

    def __init__(self,
                 certificate_authentication_api: CertificateAuthenticationApi,
                 certificate_pod_api: CertificatePodApi,
                 app_id: str,
                 retry_config: BdkRetryConfig,
                 tokens_repository: TokensRepository = None):
        """

        :param certificate_authentication_api: the CertificateAuthenticationApi instance
        :param certificate_pod_api: the CertificatePodApi instance
        :param app_id: the application ID
        :param retry_config: Retry configuration
        :param tokens_repository: the tokens repository to store existing valid sessions.
          Defaults to InMemoryTokensRepository
        """
        super().__init__(app_id, tokens_repository)
        self._certificate_authentication_api = certificate_authentication_api
        self._certificate_pod_api = certificate_pod_api
        self._retry_config = retry_config

    @retry(retry=authentication_retry)
    async def _retrieve_tokens(self, app_token: str) -> ExtensionAppTokens:
        authentication_request = ExtensionAppAuthenticateRequest(app_token=app_token)
        return await self._certificate_authentication_api.v1_authenticate_extension_app_post(
            auth_request=authentication_request)

    @retry(retry=authentication_retry)
    async def _get_pod_certificate(self) -> PodCertificate:
        return await self._certificate_pod_api.v1_app_pod_certificate_get()
