from messagebus import MessageBus, Message, Event


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

    def publish_event(self, event):
        # type: (Event) -> Message
        event.headers["context"] = self.context
        return self.message_bus.publish_event(event)

    def send_message(self, component, message):
        # type: (str, Message) -> Message
        message.headers["context"] = self.context
        return self.message_bus.send_message(component, message)

    def subscribe(self, subject):
        # type: (str) -> None
        self.message_bus.subscribe(subject)

    def reply_message(self, message, reply_message):
        # type: (Message, Message) -> None
        reply_message.headers["context"] = self.context
        self.message_bus.reply_message(message, reply_message)

    def reply_event(self, event, reply_message):
        # type: (Event, Message) -> None
        reply_message.headers["context"] = self.context
        self.message_bus.reply_event(event, reply_message)
