from messagebus import MessageBus


class ContextMessageBus(MessageBus):
    def __init__(
        self,
        context: str,
        message_bus: MessageBus,
    ) -> None:
        self.context = context
        self.message_bus = message_bus

    def publish_event(self, event) :
        event.headers["context"] = self.context
        return self.message_bus.publish_event(event)

    def send_message(self, component, message):
        message.headers["context"] = self.context
        return self.message_bus.send_message(component, message)

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

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

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