import base64
import logging
import os
import re
from email.message import EmailMessage
from typing import Any, Generator, Iterable, Tuple, Union
from urllib.parse import unquote_plus

import cachecontrol
import cryptography.hazmat.backends
import jwt
import requests
import requests.utils
from cryptography.hazmat.primitives import serialization
from requests.exceptions import ConnectionError, RequestException

from atlassian_jwt_auth.exceptions import (
    KeyIdentifierException,
    PrivateKeyRetrieverException,
    PublicKeyRetrieverException,
)

PEM_FILE_TYPE = "application/x-pem-file"


class KeyIdentifier(object):
    """This class represents a key identifier"""

    def __init__(self, identifier: str) -> None:
        self.__key_id = validate_key_identifier(identifier)

    @property
    def key_id(self) -> str:
        return self.__key_id


def validate_key_identifier(identifier: str) -> str:
    """returns a validated key identifier."""
    regex = re.compile(r"^[\w.\-\+/]*$")
    _error_msg = "Invalid key identifier %s" % identifier
    if not identifier:
        raise KeyIdentifierException(_error_msg)
    if not regex.match(identifier):
        raise KeyIdentifierException(_error_msg)
    normalised = os.path.normpath(identifier)
    if normalised != identifier:
        raise KeyIdentifierException(_error_msg)
    if normalised.startswith("/"):
        raise KeyIdentifierException(_error_msg)
    if ".." in normalised:
        raise KeyIdentifierException(_error_msg)
    return identifier


def _get_key_id_from_jwt_header(a_jwt: Union[str, bytes]) -> KeyIdentifier:
    """returns the key identifier from a jwt header."""
    header = jwt.get_unverified_header(a_jwt)
    return KeyIdentifier(header["kid"])


class BasePublicKeyRetriever(object):
    """Base class for retrieving a public key."""

    def retrieve(self, key_identifier: Union[KeyIdentifier, str], **kwargs) -> Any:
        raise NotImplementedError()


class HTTPSPublicKeyRetriever(BasePublicKeyRetriever):
    """This class retrieves public key from a https location based upon the
    given key id.
    """

    # Use a static requests session, reused/shared by all instances of
    # HTTPSPublicKeyRetriever:
    _class_session = None

    def __init__(self, base_url: str) -> None:
        if base_url is None or not base_url.startswith("https://"):
            raise PublicKeyRetrieverException("The base url must start with https://")
        if not base_url.endswith("/"):
            base_url += "/"
        self.base_url = base_url
        self._session = self._get_session()
        self._proxies = requests.utils.get_environ_proxies(self.base_url)

    def _get_session(self) -> requests.Session:
        if HTTPSPublicKeyRetriever._class_session is None:
            session = cachecontrol.CacheControl(requests.Session())
            session.trust_env = False
            HTTPSPublicKeyRetriever._class_session = session
        return HTTPSPublicKeyRetriever._class_session

    def retrieve(
        self, key_identifier: Union[KeyIdentifier, str], **requests_kwargs: Any
    ) -> Any:
        """returns the public key for given key_identifier."""
        if not isinstance(key_identifier, KeyIdentifier):
            key_identifier = KeyIdentifier(key_identifier)
        if self._proxies and "proxies" not in requests_kwargs:
            requests_kwargs["proxies"] = self._proxies
        url = self.base_url + key_identifier.key_id
        try:
            return self._retrieve(url, requests_kwargs)
        except requests.RequestException as e:
            try:
                status_code = e.response.status_code if e.response else None
            except AttributeError:
                status_code = None
            raise PublicKeyRetrieverException(e, status_code=status_code)

    def _retrieve(self, url: str, requests_kwargs: Any) -> Any:
        resp = self._session.get(
            url, headers={"accept": PEM_FILE_TYPE}, **requests_kwargs
        )
        resp.raise_for_status()
        self._check_content_type(url, resp.headers["content-type"])
        return resp.text

    def _check_content_type(self, url: str, content_type: str):
        msg = EmailMessage()
        msg["content-type"] = content_type
        media_type = msg.get_content_type()

        if media_type.lower() != PEM_FILE_TYPE.lower():
            raise PublicKeyRetrieverException(
                "Invalid content-type, '%s', for url '%s' ." % (content_type, url)
            )


class HTTPSMultiRepositoryPublicKeyRetriever(BasePublicKeyRetriever):
    """This class retrieves public key from the supplied https key
    repository locations based upon key ids.
    """

    def __init__(self, key_repository_urls: Iterable[str]) -> None:
        if not isinstance(key_repository_urls, list):
            raise TypeError("keystore_urls must be a list of urls.")
        self._retrievers = self._create_retrievers(key_repository_urls)

    def _create_retrievers(
        self, key_repository_urls: Iterable[str]
    ) -> Iterable[BasePublicKeyRetriever]:
        return [HTTPSPublicKeyRetriever(url) for url in key_repository_urls]

    def handle_retrieval_exception(
        self, retriever: BasePublicKeyRetriever, exception: Exception
    ):
        """Handles working with exceptions encountered during key
        retrieval.
        """
        if isinstance(exception, PublicKeyRetrieverException):
            original_exception = getattr(exception, "original_exception", None)
            if isinstance(original_exception, ConnectionError):
                return
            if exception.status_code is None or exception.status_code < 500:
                raise

    def retrieve(
        self, key_identifier: Union[KeyIdentifier, str], **requests_kwargs: Any
    ) -> Any:
        for retriever in self._retrievers:
            try:
                return retriever.retrieve(key_identifier, **requests_kwargs)
            except (RequestException, PublicKeyRetrieverException) as e:
                self.handle_retrieval_exception(retriever, e)
                logger = logging.getLogger(__name__)
                logger.warning(
                    "Unable to retrieve public key from store",
                    extra={
                        "underlying_error": str(e),
                        "key repository": getattr(retriever, "base_url", "unknown"),
                    },
                )
        raise PublicKeyRetrieverException("Cannot load key from key repositories")


class BasePrivateKeyRetriever(object):
    """This is the base private key retriever class."""

    def load(self, issuer: str) -> Tuple[Union[KeyIdentifier], Union[str, bytes]]:
        """returns the key identifier and private key pem found
        for the given issuer.
        """
        raise NotImplementedError("Not implemented.")


class DataUriPrivateKeyRetriever(BasePrivateKeyRetriever):
    """This class can be used to retrieve the key identifier and
    private key from the supplied data uri.
    """

    def __init__(self, data_uri: str) -> None:
        self._data_uri = data_uri

    def load(self, issuer: str) -> Tuple[Union[KeyIdentifier], Union[str, bytes]]:
        if not self._data_uri.startswith("data:application/pkcs8;kid="):
            raise PrivateKeyRetrieverException("Unrecognised data uri format.")
        splitted = self._data_uri.split(";")
        key_identifier = KeyIdentifier(unquote_plus(splitted[1][len("kid=") :]))
        key_data = base64.b64decode(splitted[-1].split(",")[-1])
        key = serialization.load_der_private_key(
            key_data,
            password=None,
            backend=cryptography.hazmat.backends.default_backend(),
        )
        private_key_pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )
        return key_identifier, private_key_pem.decode("utf-8")


class StaticPrivateKeyRetriever(BasePrivateKeyRetriever):
    """This class simply returns the key_identifier and private_key_pem
    initially provided to it in calls to load.
    """

    def __init__(
        self,
        key_identifier: Union[KeyIdentifier, str],
        private_key_pem: Union[str, bytes],
    ) -> None:
        if not isinstance(key_identifier, KeyIdentifier):
            key_identifier = KeyIdentifier(key_identifier)

        self.key_identifier: KeyIdentifier = key_identifier
        self.private_key_pem: Union[str, bytes] = private_key_pem

    def load(self, issuer: str) -> Tuple[Union[KeyIdentifier], Union[str, bytes]]:
        return self.key_identifier, self.private_key_pem


class FilePrivateKeyRetriever(BasePrivateKeyRetriever):
    """This class can be used to retrieve the latest key identifier and
    private key for a given issuer found under its private key
    repository path.
    """

    def __init__(self, private_key_repository_path: str) -> None:
        self.private_key_repository = FilePrivateKeyRepository(
            private_key_repository_path
        )

    def load(self, issuer: str) -> Tuple[KeyIdentifier, str]:
        key_identifier = self._find_last_key_id(issuer)
        private_key_pem = self.private_key_repository.load_key(key_identifier)
        return key_identifier, private_key_pem

    def _find_last_key_id(self, issuer) -> KeyIdentifier:
        key_identifiers = list(self.private_key_repository.find_valid_key_ids(issuer))

        if key_identifiers:
            return key_identifiers[-1]
        else:
            raise IOError("Issuer has no valid keys: %s" % issuer)


class FilePrivateKeyRepository(object):
    """This class represents a file backed private key repository."""

    def __init__(self, path) -> None:
        self.path = path

    def find_valid_key_ids(self, issuer: str) -> Generator[KeyIdentifier, Any, None]:
        issuer_directory = os.path.join(self.path, issuer)
        for filename in sorted(os.listdir(issuer_directory)):
            if filename.endswith(".pem"):
                yield KeyIdentifier("%s/%s" % (issuer, filename))

    def load_key(self, key_identifier: KeyIdentifier) -> str:
        key_filename = os.path.join(self.path, key_identifier.key_id)
        with open(key_filename, "rb") as f:
            return f.read().decode("utf-8")
