import logging

from defence360agent.contracts.config import SystemConfig
from defence360agent.contracts.messages import MessageType
from defence360agent.contracts.plugins import MessageSink, expect
from defence360agent.subsys import svcctl
from im360.contracts.config import Webshield

logger = logging.getLogger(__name__)


class WebshieldFeatureFlagsPlugin(MessageSink):
    def __init__(self):
        self._loop = None
        self._use_feature_flags = Webshield.USE_FEATURE_FLAGS

    async def create_sink(self, loop) -> None:
        self._loop = loop

    @expect(MessageType.ConfigUpdate)
    async def on_config_update(self, message: MessageType.ConfigUpdate):
        if not isinstance(message["conf"], SystemConfig):
            return
        current = Webshield.USE_FEATURE_FLAGS
        if current == self._use_feature_flags:
            return
        logger.info(
            "WEBSHIELD.use_feature_flags changed from %s to %s,"
            " restarting agent",
            self._use_feature_flags,
            current,
        )
        self._use_feature_flags = current
        await svcctl.imunify360_agent_service().restart()