from judy_core.messages import (
    EnvironmentNotification,
    EnvironmentNotificationList,
    GetEnvironmentNotification,
)
from tongcell.core import MessageContext, RoutedActor, message_handler


class NotificationCenter(RoutedActor):
    def __init__(self):
        super().__init__()
        self._notifications = EnvironmentNotificationList()

    @message_handler()
    async def stage_notification(
        self, notification: EnvironmentNotification, ctx: MessageContext
    ) -> None:
        self._notifications.notifications.append(notification)

    @message_handler()
    async def get_notifications(
        self, msg: GetEnvironmentNotification
    ) -> EnvironmentNotificationList:
        notifications, self._notifications = (
            self._notifications,
            EnvironmentNotificationList(),
        )
        return notifications
