import asyncio
import base64
import copy
import json
import time
import uuid
from typing import Any, Awaitable, Callable, Dict, List

import aiofiles
import aiohttp
import jwt

from .. import app_config, util
from ..broker import Broker, BrokerOptions
from ..logger import Logger
from ..request import Request
from ..response import Response

NAME = 'RestHTTP'
"""Name of the main class of the module.

Used by the App class
to obtain the correct class via getattr

"""


class OauthOptions:
    """Options specific to JWT/oauth authorization.

    Attributes:
        service_key_file (str): Path to a JSON service account key file to use
            for JWT generation and Bearer authorization
        scope (str): JWT scope to request authorization for
        uri (str): oauth request URI
        algorithm (str): RS256 etc
    """
    service_key_file: str
    scope: str
    uri: str
    algorithm: str

    def unique_id(self) -> str:
        """Returns a unique string that can be used in a dict."""
        return f'service file {self.service_key_file} scope {self.scope}'


class JWTAuth:
    """Container to manage JWT tokens.

    Attributes:
        _logger (Logger): logger instance
        _oauth (OauthOptions): information needed to create/generate tokens
        _service_key (dict): JSON-decoded service key
        _expires (int): expiry time of current access token
        _token (str): current access token if any
    """
    _logger: Logger
    _oauth: OauthOptions
    _service_key: Dict[str, str] = {}
    _expires: int = -1
    _token: str = ''

    def __init__(self, logger: Logger, oauth: OauthOptions):
        """Constructor for JWTAuth object.

        Args:
            logger (Logger): logger instance
            oauth (OauthOptions): information needed to create/generate tokens
        """
        self._logger = logger
        self._oauth = oauth

    async def _load_service_key(self) -> None:
        # TODO: check file modification time and reload if needed
        if 'client_email' in self._service_key:
            return

        async with aiofiles.open(self._oauth.service_key_file) as fh:
            self._service_key = json.loads(await fh.read())

    async def token(self) -> str:
        """Returns a current authorization token, refreshing it if needed."""
        now = time.time()

        cutoff = now - 10
        if self._token and self._expires >= 0 and self._expires > cutoff:
            return self._token

        await self._load_service_key()

        expiry = now + 3600  # max duration
        self._token = ''  # invalidate old token
        self._expires = -1

        encoded_jwt = jwt.encode(
            {
                'iss': self._service_key['client_email'],
                'scope': self._oauth.scope,
                'aud': self._service_key['token_uri'],
                'iat': now,
                'exp': expiry,
            },
            self._service_key['private_key'],
            algorithm=self._oauth.algorithm,
        )

        async with aiohttp.ClientSession() as session:
            self._logger.info('Refreshing JWT/oauth token for ' +
                              f'{self._oauth.service_key_file}')

            data = {
                'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',
                'assertion': encoded_jwt,
            }

            headers: Dict[str, str] = {}

            async with session.request(
                'POST', self._oauth.uri, headers=headers, data=data,
            ) as response:
                print('Status:', response.status)
                print('Content-type:', response.headers['content-type'])

                content = await response.text()
                print('Body:', content)
                bb = await response.json()
                print(type(bb))
                print(bb)

                print(bb['access_token'])
                print(bb['expires_in'])
                bearer = bb['access_token']
                self._token = f'Bearer {bearer}'
                self._expires = now + bb['expires_in']

        return self._token


class RestHTTPOptions(BrokerOptions):
    """Options for the RestHTTP broker.

    Attributes:
        method (str): GET or POST
        uri (str): HTTP or HTTPS URI. May contain macros
        headers (dict): Additional headers to add
        oauth (OauthOptions): Optional oauth/JWT config
        body_type (str): Either "urlencoded" (default) for a form POST,
            "json" for JSON encoded payload, "binary" for raw binary POST,
            or "form-data" for a MIME POST
        body (dict): key/value pairs for request body
            ${data} as value will be replaced by the data from the request
            ${data.binary} as value will be replace with the base64-decoded
                data
        response (dict): Optional format for response data. If not set, the raw
            response body will be returned. Only applicable for JSON responses
        wait (bool): Optionally enable "wait for complete" operation
        sleep (int): With wait=True, sleep this many seconds between requests,
            waiting for completion
        timeout (int): Maximum time to wait for completion before returning
            error
    """
    method: str
    uri: str
    headers: Dict[str, str] = {}
    oauth: OauthOptions
    body_type: str = 'urlencoded'
    body: Dict[str, Any] = {}
    response: Dict[str, Any] = {}
    wait: bool = False
    sleep: int = 1
    timeout: int = 600


class RestHTTP(Broker):
    """RestHTTP broker class, inherits Broker.

    Attributes:
        _jwt_cache (dict): internally managed cache of JWT tokens
    """
    _jwt_cache: Dict[str, JWTAuth] = {}

    def __init__(self, send_req_cb: Callable[[Request],
                                             Awaitable[None]]) -> None:
        """Constructor for RestHTTP class.

        Args:
            send_req_cb (callable): callback that a broker can use to trigger
                sending another request. Set by the App class.

        Returns:
            None

        """
        self._logger = Logger(util.class_name(self))
        super().__init__(send_req_cb)

    async def start(self) -> None:
        """Starts the broker.

        Currently does nothing as there is no
        specific logic related to the broker startup. Reserved for
        future use and consistency with the parent class

        Args:
            None

        Returns
            None

        """
        pass

    async def _get_jwt(self, oauth: OauthOptions) -> str:
        unique = oauth.unique_id()
        if unique not in self._jwt_cache:
            self._jwt_cache[unique] = JWTAuth(self._logger, oauth)
        jwt_auth = self._jwt_cache[unique]
        return await jwt_auth.token()

    def _apply_response_format(self,
                               response: Any,
                               extra_dict: Dict[str, Any],
                               http_options: RestHTTPOptions) -> Any:
        if not isinstance(response, dict):
            return response
        if not http_options.response:
            return response

        macros = extra_dict | response
        new_response = copy.deepcopy(http_options.response)
        util.apply_dict_macros(new_response, macros)
        return new_response

    async def process_request(self,
                              request: Request,
                              options: Dict[str, Any],
                              callback: Callable[[Request, Response],
                                                 Awaitable[None]]) -> None:
        """Processes the provided request and sends back responses.

        Uses the provided callback function

        Args:
            request (Request): Request message object
            options (Dict): broker options
            callback (Callable[[Request, Response], Awaitable[None]]):
                callback function to dispatch the responses to

        Returns
            None

        """
        http_options = RestHTTPOptions()
        util.apply_dict(options, http_options)

        response = Response()
        response.status = 'accepted'
        await asyncio.ensure_future(callback(request, response))

        headers = copy.deepcopy(http_options.headers)

        uri_macros: Dict[str, Any] = {}
        uri = http_options.uri

        if isinstance(request.data, dict):
            uri_macros = request.data | uri_macros
        elif isinstance(request.data, str):
            uri_macros['data'] = request.data

        # generate new UUID if required
        if uri.find('${uuid}') != -1 \
                and 'uuid' not in uri_macros:
            uri_macros['uuid'] = str(uuid.uuid4())

        uri = util.apply_str_macros(uri, uri_macros)

        response = Response()

        do_wait: bool = True
        start: float = time.time()

        while do_wait:
            if not http_options.wait:
                do_wait = False  # single shot
            else:
                await asyncio.sleep(http_options.sleep)

                if start + http_options.timeout < time.time():
                    do_wait = False  # last try

            try:
                if 'oauth' in http_options.__dict__ \
                        and 'scope' in http_options.oauth.__dict__ \
                        and http_options.oauth.scope != '':
                    auth = await self._get_jwt(http_options.oauth)
                    if not auth:
                        raise NameError('no authorization')
                    headers['Authorization'] = auth

                async with aiohttp.ClientSession() as session:
                    body = copy.deepcopy(http_options.body)
                    req_body: str | Dict[str, Any] | bytes | aiohttp.FormData

                    if http_options.body_type == 'binary':
                        req_body = base64.b64decode(request.data)
                    elif http_options.body_type == 'form-data':
                        util.apply_dict_macros(body, request.data)
                        req_body = aiohttp.FormData()
                        for key, val in body.items():
                            if isinstance(val, dict):
                                req_body.add_field(key, val['content'],
                                                   content_type=val.get(
                                                       'content_type'
                                                       ))
                            else:
                                req_body.add_field(key, val)
                    else:
                        util.apply_dict_macros(body, request.data)
                        req_body = body
                        if http_options.body_type == 'json':
                            req_body = json.dumps(body)

                    self.log.info(f'Making HTTP {http_options.method} ' +
                                  f'request to {uri}')
                    self.log.debug(util.compact_str(f'HTTP request body: \
                            {str(req_body)}'))

                    async with session.request(
                            http_options.method,
                            uri,
                            headers=headers,
                            data=req_body) as http_response:
                        if http_response.status < 200 \
                                or http_response.status >= 300:
                            resp_body = await http_response.text()
                            self.log.error(util.compact_str(f'Got HTTP status \
                                    {http_response.status}: {resp_body}'))
                            response.status = 'error'
                            response.reason = util.compact_str(f'status code \
                                    {http_response.status}: {resp_body}')
                        else:
                            ctype = http_response.headers['content-type']
                            self.log.debug(f'HTTP response type {ctype}')
                            if ctype == 'application/json' \
                                    or ctype.startswith('application/json ') \
                                    or ctype.startswith('application/json;'):
                                response.data = await http_response.json()
                            else:
                                response.data = await http_response.text()

                            self.log.debug('HTTP response body ' +
                                           str(response.data))
                            response.status = 'done'

                            # check for done/wait loop
                            if http_options.wait \
                                    and isinstance(response.data, dict):
                                if 'done' in response.data \
                                        and response.data['done']:
                                    do_wait = False
                                else:
                                    self.log.debug('Response not ready yet')

                            # apply response format if we're done
                            if not do_wait and isinstance(response.data, dict):
                                response.data = self._apply_response_format(
                                        response.data,
                                        uri_macros,
                                        http_options)

            except (aiohttp.client_exceptions.ClientError, NameError) as err:
                self.log.debug(util.compact_str(f'HTTP client exception: \
                        {str(err)}'))
                response.reason = f'exception: {str(err)}'
                response.status = 'error'

        await super().follow_tasks(request, response, http_options)

        await asyncio.ensure_future(callback(request, response))
