import threading
from typing import Dict

from messagebus import MessageBus


class MultiContextMessageBus(MessageBus):
    def __init__(self, message_bus, contexts):
        # type: (MessageBus, Dict[str, str]) -> None
        self.message_bus = message_bus
        self.contexts = contexts
        self.local_data = threading.local()

    def publish_event(self, event):
        try:
            context = self.local_data.context
            event.headers["context"] = context
        except AttributeError:
            pass
        return self.message_bus.publish_event(event)

    def send_message(self, component, message):
        try:
            context = self.local_data.context
            message.headers["context"] = context
        except AttributeError:
            pass

        return self.message_bus.send_message(component, message)

    def subscribe(self, subscriptions):
        self.message_bus.subscribe(subscriptions)

    def reply_message(self, message, reply_message):
        try:
            context = reply_message.headers["context"]
            message.headers["context"] = context
            self.message_bus.reply_message(message, reply_message)
        except KeyError:
            self.message_bus.reply_message(message, reply_message)

    def reply_event(self, event, reply_message):
        try:
            context = reply_message.headers["context"]
            event.headers["context"] = context
            self.message_bus.reply_event(event, reply_message)
        except KeyError:
            self.message_bus.reply_event(event, reply_message)

    def set_message_handler(self, message_handler):
        self.message_bus.set_message_handler(message_handler)

    def get_reply(self):
        return self.message_bus.get_reply()

    def set_current_context(self, context):
        self.local_data.context = context
