import asyncio
import copy
import time
from collections.abc import Iterable

import agentspeak
import agentspeak.runtime
import agentspeak.stdlib
from pydantic import ValidationError

from control_backend.agents.base import BaseAgent
from control_backend.core.agent_system import InternalMessage
from control_backend.core.config import settings
from control_backend.schemas.belief_message import Belief, BeliefMessage
from control_backend.schemas.llm_prompt_message import LLMPromptMessage
from control_backend.schemas.ri_message import SpeechCommand

DELIMITER = ";\n"  # TODO: temporary until we support lists in AgentSpeak


class BDICoreAgent(BaseAgent):
    """
    BDI Core Agent.

    This is the central reasoning agent of the system, powered by the **AgentSpeak(L)** language.
    It maintains a belief base (representing the state of the world) and a set of plans (rules).

    It runs an internal BDI (Belief-Desire-Intention) cycle using the ``agentspeak`` library.
    When beliefs change (e.g., via :meth:`_apply_beliefs`), the agent evaluates its plans to
    determine the best course of action.

    **Custom Actions:**
    It defines custom actions (like ``.reply``) that allow the AgentSpeak code to interact with
    external Python agents (e.g., querying the LLM).

    :ivar bdi_agent: The internal AgentSpeak agent instance.
    :ivar asl_file: Path to the AgentSpeak source file (.asl).
    :ivar env: The AgentSpeak environment.
    :ivar actions: A registry of custom actions available to the AgentSpeak code.
    :ivar _wake_bdi_loop: Event used to wake up the reasoning loop when new beliefs arrive.
    """

    bdi_agent: agentspeak.runtime.Agent

    def __init__(self, name: str, asl: str):
        super().__init__(name)
        self.asl_file = asl
        self.env = agentspeak.runtime.Environment()
        # Deep copy because we don't actually want to modify the standard actions globally
        self.actions = copy.deepcopy(agentspeak.stdlib.actions)
        self._wake_bdi_loop = asyncio.Event()

    async def setup(self) -> None:
        """
        Initialize the BDI agent.

        1. Registers custom actions (like ``.reply``).
        2. Loads the .asl source file.
        3. Starts the reasoning loop (:meth:`_bdi_loop`) in the background.
        """
        self.logger.debug("Setup started.")

        self._add_custom_actions()

        await self._load_asl()

        # Start the BDI cycle loop
        self.add_behavior(self._bdi_loop())
        self._wake_bdi_loop.set()
        self.logger.debug("Setup complete.")

    async def _load_asl(self):
        """
        Load and parse the AgentSpeak source file.
        """
        try:
            with open(self.asl_file) as source:
                self.bdi_agent = self.env.build_agent(source, self.actions)
        except FileNotFoundError:
            self.logger.warning(f"Could not find the specified ASL file at {self.asl_file}.")
            self.bdi_agent = agentspeak.runtime.Agent(self.env, self.name)

    async def _bdi_loop(self):
        """
        The main BDI reasoning loop.

        It waits for the ``_wake_bdi_loop`` event (set when beliefs change or actions complete).
        When awake, it steps through the AgentSpeak interpreter. It also handles sleeping if
        the agent has deferred intentions (deadlines).
        """
        while self._running:
            await (
                self._wake_bdi_loop.wait()
            )  # gets set whenever there's an update to the belief base

            # Agent knows when it's expected to have to do its next thing
            maybe_more_work = True
            while maybe_more_work:
                maybe_more_work = False
                self.logger.debug("Stepping BDI.")
                if self.bdi_agent.step():
                    maybe_more_work = True

                if not maybe_more_work:
                    deadline = self.bdi_agent.shortest_deadline()
                    if deadline:
                        self.logger.debug("Sleeping until %s", deadline)
                        await asyncio.sleep(deadline - time.time())
                        maybe_more_work = True
                    else:
                        self._wake_bdi_loop.clear()
                        self.logger.debug("No more deadlines. Halting BDI loop.")

    async def handle_message(self, msg: InternalMessage):
        """
        Handle incoming messages.

        - **Beliefs**: Updates the internal belief base.
        - **LLM Responses**: Forwards the generated text to the Robot Speech Agent (actuation).

        :param msg: The received internal message.
        """
        self.logger.debug("Processing message from %s.", msg.sender)

        if msg.thread == "beliefs":
            try:
                beliefs = BeliefMessage.model_validate_json(msg.body).beliefs
                self._apply_beliefs(beliefs)
            except ValidationError:
                self.logger.exception("Error processing belief.")
            return

        # The message was not a belief, handle special cases based on sender
        match msg.sender:
            case settings.agent_settings.llm_name:
                content = msg.body
                self.logger.info("Received LLM response: %s", content)

                # Forward to Robot Speech Agent
                cmd = SpeechCommand(data=content)
                out_msg = InternalMessage(
                    to=settings.agent_settings.robot_speech_name,
                    sender=self.name,
                    body=cmd.model_dump_json(),
                )
                await self.send(out_msg)

    def _apply_beliefs(self, beliefs: list[Belief]):
        """
        Update the belief base with a list of new beliefs.

        If ``replace=True`` is set on a belief, it removes all existing beliefs with that name
        before adding the new one.
        """
        if not beliefs:
            return

        for belief in beliefs:
            if belief.replace:
                self._remove_all_with_name(belief.name)
            self._add_belief(belief.name, belief.arguments)

    def _add_belief(self, name: str, args: Iterable[str] = []):
        """
        Add a single belief to the BDI agent.

        :param name: The functor/name of the belief (e.g., "user_said").
        :param args: Arguments for the belief.
        """
        # new_args = (agentspeak.Literal(arg) for arg in args)  # TODO: Eventually support multiple
        merged_args = DELIMITER.join(arg for arg in args)
        new_args = (agentspeak.Literal(merged_args),)
        term = agentspeak.Literal(name, new_args)

        self.bdi_agent.call(
            agentspeak.Trigger.addition,
            agentspeak.GoalType.belief,
            term,
            agentspeak.runtime.Intention(),
        )

        self._wake_bdi_loop.set()

        self.logger.debug(f"Added belief {self.format_belief_string(name, args)}")

    def _remove_belief(self, name: str, args: Iterable[str]):
        """
        Removes a specific belief (with arguments), if it exists.
        """
        new_args = (agentspeak.Literal(arg) for arg in args)
        term = agentspeak.Literal(name, new_args)

        result = self.bdi_agent.call(
            agentspeak.Trigger.removal,
            agentspeak.GoalType.belief,
            term,
            agentspeak.runtime.Intention(),
        )

        if result:
            self.logger.debug(f"Removed belief {self.format_belief_string(name, args)}")
            self._wake_bdi_loop.set()
        else:
            self.logger.debug("Failed to remove belief (it was not in the belief base).")

    def _remove_all_with_name(self, name: str):
        """
        Removes all beliefs that match the given `name`.
        """
        relevant_groups = []
        for key in self.bdi_agent.beliefs:
            if key[0] == name:
                relevant_groups.append(key)

        removed_count = 0
        for group in relevant_groups:
            beliefs_to_remove = list(self.bdi_agent.beliefs[group])
            for belief in beliefs_to_remove:
                self.bdi_agent.call(
                    agentspeak.Trigger.removal,
                    agentspeak.GoalType.belief,
                    belief,
                    agentspeak.runtime.Intention(),
                )
                removed_count += 1

        self._wake_bdi_loop.set()

        self.logger.debug(f"Removed {removed_count} beliefs.")

    def _add_custom_actions(self) -> None:
        """
        Add any custom actions here. Inside `@self.actions.add()`, the first argument is
        the name of the function in the ASL file, and the second the amount of arguments
        the function expects (which will be located in `term.args`).
        """

        @self.actions.add(".reply", 3)
        def _reply(agent: "BDICoreAgent", term, intention):
            """
            Sends text to the LLM (AgentSpeak action).
            Example: .reply("Hello LLM!", "Some norm", "Some goal")
            """
            message_text = agentspeak.grounded(term.args[0], intention.scope)
            norms = agentspeak.grounded(term.args[1], intention.scope)
            goals = agentspeak.grounded(term.args[2], intention.scope)

            self.logger.debug("Norms: %s", norms)
            self.logger.debug("Goals: %s", goals)
            self.logger.debug("User text: %s", message_text)

            asyncio.create_task(self._send_to_llm(str(message_text), str(norms), str(goals)))
            yield

    async def _send_to_llm(self, text: str, norms: str = None, goals: str = None):
        """
        Sends a text query to the LLM agent asynchronously.
        """
        prompt = LLMPromptMessage(
            text=text,
            norms=norms.split("\n") if norms else [],
            goals=goals.split("\n") if norms else [],
        )
        msg = InternalMessage(
            to=settings.agent_settings.llm_name,
            sender=self.name,
            body=prompt.model_dump_json(),
        )
        await self.send(msg)
        self.logger.info("Message sent to LLM agent: %s", text)

    @staticmethod
    def format_belief_string(name: str, args: Iterable[str] = []):
        """
        Given a belief's name and its args, return a string of the form "name(*args)"
        """
        return f"{name}{'(' if args else ''}{','.join(args)}{')' if args else ''}"
