from datetime import datetime
from typing import Any
import base64
import json
from itertools import count

from websocket import WebSocket
from demessaging.PulsarMessageConstants import PulsarConfigKeys, MessageType, PropertyKeys
from demessaging.PulsarConnection import PulsarConnection


class PulsarMessageProducer(PulsarConnection):
    """Producer class to send requests to a registered backend module (topic)"""

    def __init__(self, pulsar_config):
        super().__init__(pulsar_config)
        self.subscription_name: str = 'python-backend-' + datetime.now().isoformat()
        self.context_counter = count()

    async def send_request(self, request_msg, topic: str = None) -> Any:
        """Sends the given request to the backend module bound to the topic provided in the pulsar configuration.
        In order to increase re-usability the destination topic can be overridden with the optional topic argument.

        :param request_msg: dictionary providing a 'property' dictionary, a payload string, or both
        :param topic: overrides the used topic for this request
        :return: received response from the backend module
        """
        # establish connections
        out_socket: WebSocket = self.open_socket(topic=topic)
        # To be thread safe, we generate the response topic here
        response_topic: str = self.generate_response_topic()
        in_socket: WebSocket = self.open_socket(subscription=self.subscription_name, topic=response_topic)

        try:
            # create message context (from counter)
            context = next(self.context_counter)
            request_msg['context'] = context
            if 'properties' not in request_msg:
                request_msg['properties'] = {}
            request_msg['properties'][PropertyKeys.RESPONSE_TOPIC] = response_topic
            request_msg['properties'][PropertyKeys.REQUEST_CONTEXT] = request_msg['context']
            request_msg['properties'][PropertyKeys.MESSAGE_TYPE] = MessageType.REQUEST

            # send message via outgoing connection to request topic
            out_socket.send(json.dumps(request_msg))
            ack = out_socket.recv()
            ack = json.loads(ack)
            if 'error' in ack['result']:
                # error - return
                return {'status': 'error',
                        'error': 'error sending the request',
                        'msg': ack}

            # wait for response on response topic
            response = in_socket.recv()

            # parse json message
            response = json.loads(response)
            props = response['properties']

            # assert that we received a 'response' message and check for matching context
            if props[PropertyKeys.MESSAGE_TYPE] != MessageType.RESPONSE:
                return {'status': 'error',
                        'error': 'received message is not a response to the sent request',
                        'msg': response}

            if int(props[PropertyKeys.REQUEST_CONTEXT]) != int(context):
                return {'status': 'error',
                        'error': 'received message belongs to a different request',
                        'msg': response}

            if 'payload' in response:
                status = 'success'
                if 'status' in props:
                    # we might successfully get a response, but it might contain an error from the backend
                    status = props['status']

                # decode b64 payload msg
                payload: str = response['payload']

                try:
                    payload = base64.b64decode(payload).decode("utf-8")
                except Exception as e:
                    status = 'error'
                    payload = 'error decoding payload: {0}'.format(e)

                if status == 'error':
                    return {'status': 'error', 'error': payload, 'msg': response}
                else:
                    return {'status': status, 'msg': payload}
            else:
                return {'status': 'error',
                        'error': 'missing response payload',
                        'msg': response}
        finally:
            # close the sockets
            if out_socket:
                out_socket.close()
            if in_socket:
                in_socket.close()
