from typing import Optional, Union

from asn1crypto import cms, tsp
from pyhanko_certvalidator.fetchers.aiohttp_fetchers.util import LazySession

from pyhanko.sign.timestamps import TimeStamper, TimestampRequestError
from pyhanko.sign.timestamps.common_utils import set_tsp_headers

try:
    import aiohttp
except ImportError as _e:  # pragma: nocover
    raise ImportError(
        "Install pyHanko with the [async_http] optional dependency group"
    ) from _e


class AIOHttpTimeStamper(TimeStamper):
    def __init__(self, url, session: Union[aiohttp.ClientSession, LazySession],
                 https=False, timeout=5, headers=None,
                 auth: Optional[aiohttp.BasicAuth] = None):
        """
        Initialise the timestamp client.

        :param url:
            URL where the server listens for timestamp requests.
        :param https:
            Enforce HTTPS.
        :param timeout:
            Timeout (in seconds)
        :param auth:
            `aiohttp.BasicAuth` object with authentication credentials.
        :param headers:
            Other headers to include.
        """
        if https and not url.startswith('https:'):  # pragma: nocover
            raise ValueError('Timestamp URL is not HTTPS.')
        self.url = url
        self.timeout = timeout
        self.auth = auth
        self.headers = headers
        self._session = session
        super().__init__()

    async def async_request_headers(self) -> dict:
        """
        Format the HTTP request headers.
        Subclasses can override this to perform their own header generation
        logic.

        :return:
            Header dictionary.
        """
        return set_tsp_headers(self.headers or {})

    async def get_session(self) -> aiohttp.ClientSession:
        session = self._session
        if isinstance(session, LazySession):
            return await session.get_session()
        else:
            return session

    async def async_timestamp(self, message_digest, md_algorithm) \
            -> cms.ContentInfo:
        return await super().async_timestamp(message_digest, md_algorithm)

    async def async_request_tsa_response(self, req: tsp.TimeStampReq) \
            -> tsp.TimeStampResp:
        session = await self.get_session()

        cl_timeout = aiohttp.ClientTimeout(total=self.timeout)
        headers = await self.async_request_headers()
        try:
            async with session.post(url=self.url, headers=headers,
                                    data=req.dump(), auth=self.auth,
                                    raise_for_status=True,
                                    timeout=cl_timeout) as response:
                response_data = await response.read()
                ct = response.headers.get('Content-Type')
                if ct != 'application/timestamp-reply':
                    msg = (
                        f'Bad content type. Expected '
                        f'application/timestamp-reply,but got {ct}.'
                    )
                    raise aiohttp.ContentTypeError(
                        response.request_info, response.history,
                        message=msg, headers=response.headers,
                    )
        except aiohttp.ClientError as e:
            raise TimestampRequestError(
                "Error while contacting timestamp service",
            ) from e
        return tsp.TimeStampResp.load(response_data)