import json

from spade.agent import Message
from spade.behaviour import CyclicBehaviour
from spade_bdi.bdi import BDIAgent

from control_backend.core.config import settings


class BeliefSetterBehaviour(CyclicBehaviour):
    """
    This is the behaviour that the BDI agent runs. This behaviour waits for incoming
    message and updates the agent's beliefs accordingly.
    """

    agent: BDIAgent

    async def run(self):
        """Polls for messages and processes them."""
        msg = await self.receive()
        self.agent.logger.debug(
            "Received message from %s with thread '%s' and body: %s",
            msg.sender,
            msg.thread,
            msg.body,
        )
        self._process_message(msg)

    def _process_message(self, message: Message):
        """Routes the message to the correct processing function based on the sender."""
        sender = message.sender.node  # removes host from jid and converts to str
        self.agent.logger.debug("Processing message from sender: %s", sender)

        match sender:
            case settings.agent_settings.belief_collector_agent_name:
                self.agent.logger.debug(
                    "Message is from the belief collector agent. Processing as belief message."
                )
                self._process_belief_message(message)
            case _:
                self.agent.logger.debug("Not the belief agent, discarding message")
                pass

    def _process_belief_message(self, message: Message):
        if not message.body:
            self.agent.logger.debug("Ignoring message with empty body from %s", message.sender.node)
            return

        match message.thread:
            case "beliefs":
                try:
                    beliefs: dict[str, list[str]] = json.loads(message.body)
                    self._set_beliefs(beliefs)
                except json.JSONDecodeError:
                    self.agent.logger.error(
                        "Could not decode beliefs from JSON. Message body: '%s'",
                        message.body,
                        exc_info=True,
                    )
            case _:
                pass

    def _set_beliefs(self, beliefs: dict[str, list[str]]):
        """Removes previous values for beliefs and updates them with the provided values."""
        if self.agent.bdi is None:
            self.agent.logger.warning("Cannot set beliefs; agent's BDI is not yet initialized.")
            return

        if not beliefs:
            self.agent.logger.debug("Received an empty set of beliefs. No beliefs were updated.")
            return

        # Set new beliefs (outdated beliefs are automatically removed)
        for belief, arguments in beliefs.items():
            self.agent.logger.debug("Setting belief %s with arguments %s", belief, arguments)
            self.agent.bdi.set_belief(belief, *arguments)

            # Special case: if there's a new user message, flag that we haven't responded yet
            if belief == "user_said":
                self.agent.bdi.set_belief("new_message")
                self.agent.logger.debug(
                    "Detected 'user_said' belief, also setting 'new_message' belief."
                )

        self.agent.logger.info("Successfully updated %d beliefs.", len(beliefs))