"""
This program has been developed by students from the bachelor Computer Science at Utrecht
University within the Software Project course.
© Copyright Utrecht University (Department of Information and Computing Sciences)
"""

import asyncio
import copy
import json
import logging
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 BeliefMessage
from control_backend.schemas.llm_prompt_message import LLMPromptMessage
from control_backend.schemas.ri_message import GestureCommand, RIEndpoint, SpeechCommand

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


experiment_logger = logging.getLogger(settings.logging_settings.experiment_logger_name)


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):
        super().__init__(name)
        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()
        self._bdi_loop_task = None

    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._bdi_loop_task = self.add_behavior(self._bdi_loop())
        self._wake_bdi_loop.set()
        self.logger.debug("Setup complete.")

    async def _load_asl(self, file_name: str | None = None) -> None:
        """
        Load and parse the AgentSpeak source file.
        """
        file_name = file_name or "src/control_backend/agents/bdi/default_behavior.asl"

        try:
            with open(file_name) as source:
                self.bdi_agent = self.env.build_agent(source, self.actions)
                self.logger.info(f"Loaded new ASL from {file_name}.")
        except FileNotFoundError:
            self.logger.warning(f"Could not find the specified ASL file at {file_name}.")
            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
                if self.bdi_agent.step():
                    maybe_more_work = True

                if not maybe_more_work:
                    deadline = self.bdi_agent.shortest_deadline()
                    if 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.
        - **Program**: Updates the internal agentspeak file to match the current program.
        - **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:
                belief_changes = BeliefMessage.model_validate_json(msg.body)
                self._apply_belief_changes(belief_changes)
            except ValidationError:
                self.logger.exception("Error processing belief.")
            return

        # New agentspeak file
        if msg.thread == "new_program":
            if self._bdi_loop_task:
                self._bdi_loop_task.cancel()
            await self._load_asl(msg.body)
            self.add_behavior(self._bdi_loop())

        # 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)
            case settings.agent_settings.user_interrupt_name:
                self.logger.debug("Received user interruption: %s", msg)

                match msg.thread:
                    case "force_phase_transition":
                        self._set_goal("transition_phase")
                    case "force_trigger":
                        self._force_trigger(msg.body)
                    case "force_norm":
                        self._force_norm(msg.body)
                    case "force_next_phase":
                        self._force_next_phase()
                    case "stop":
                        self._stop()
                    case _:
                        self.logger.warning("Received unknown user interruption: %s", msg)

    def _apply_belief_changes(self, belief_changes: BeliefMessage):
        """
        Update the belief base with a list of new beliefs.

        For beliefs in ``belief_changes.replace``, it removes all existing beliefs with that name
        before adding one new one.

        :param belief_changes: The changes in beliefs to apply.
        """
        if not belief_changes.create and not belief_changes.replace and not belief_changes.delete:
            return

        for belief in belief_changes.create:
            self._add_belief(belief.name, belief.arguments)

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

        for belief in belief_changes.delete:
            self._remove_belief(belief.name, belief.arguments)

    def _add_belief(self, name: str, args: list[str] = None):
        """
        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
        args = args or []
        if args:
            merged_args = DELIMITER.join(arg for arg in args)
            new_args = (agentspeak.Literal(merged_args),)
            term = agentspeak.Literal(name, new_args)
        else:
            term = agentspeak.Literal(name)

        if name != "user_said":
            experiment_logger.observation(f"Formed new belief: {name}{f'={args}' if args else ''}")

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

        # Check for transitions
        self.bdi_agent.call(
            agentspeak.Trigger.addition,
            agentspeak.GoalType.achievement,
            agentspeak.Literal("transition_phase"),
            agentspeak.runtime.Intention(),
        )

        # Check triggers
        self.bdi_agent.call(
            agentspeak.Trigger.addition,
            agentspeak.GoalType.achievement,
            agentspeak.Literal("check_triggers"),
            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] | None):
        """
        Removes a specific belief (with arguments), if it exists.
        """
        if args is None:
            term = agentspeak.Literal(name)
        else:
            new_args = (agentspeak.Literal(arg) for arg in args)
            term = agentspeak.Literal(name, new_args)

        if name != "user_said":
            experiment_logger.observation(f"Removed belief: {name}{f'={args}' if args else ''}")

        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 _set_goal(self, name: str, args: Iterable[str] | None = None):
        args = args or []

        if args:
            merged_args = DELIMITER.join(arg for arg in args)
            new_args = (agentspeak.Literal(merged_args),)
            term = agentspeak.Literal(name, new_args)
        else:
            term = agentspeak.Literal(name)

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

        self._wake_bdi_loop.set()

        self.logger.debug(f"Set goal !{self.format_belief_string(name, args)}.")

    def _force_trigger(self, name: str):
        self._set_goal(name)

        self.logger.info("Manually forced trigger %s.", name)

    # TODO: make this compatible for critical norms
    def _force_norm(self, name: str):
        self._add_belief(f"force_{name}")

        self.logger.info("Manually forced norm %s.", name)

    def _force_next_phase(self):
        self._set_goal("force_transition_phase")

        self.logger.info("Manually forced phase transition.")

    def _stop(self):
        self._set_goal("stop")

        self.logger.info("Stopped the program (skipped to end phase).")

    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(".substring_case_insensitive", 3)
        @agentspeak.optimizer.function_like
        def _substring(agent, term, intention):
            """
            Find out if a string is a substring of another (case insensitive). Copied mostly from
            the agentspeak library method .substring.
            """
            needle = agentspeak.asl_str(agentspeak.grounded(term.args[0], intention.scope)).lower()
            haystack = agentspeak.asl_str(agentspeak.grounded(term.args[1], intention.scope)).lower()

            choicepoint = object()

            pos = haystack.find(needle)
            while pos != -1:
                intention.stack.append(choicepoint)

                if agentspeak.unify(term.args[2], pos, intention.scope, intention.stack):
                    yield

                agentspeak.reroll(intention.scope, intention.stack, choicepoint)
                pos = haystack.find(needle, pos + 1)

        @self.actions.add(".reply", 2)
        def _reply(agent, term, intention):
            """
            Let the LLM generate a response to a user's utterance with the current norms and goals.
            """
            message_text = agentspeak.grounded(term.args[0], intention.scope)
            norms = agentspeak.grounded(term.args[1], intention.scope)

            self.add_behavior(self._send_to_llm(str(message_text), str(norms), ""))
            yield

        @self.actions.add(".reply_with_goal", 3)
        def _reply_with_goal(agent, term, intention):
            """
            Let the LLM generate a response to a user's utterance with the current norms and a
            specific goal.
            """
            message_text = agentspeak.grounded(term.args[0], intention.scope)
            norms = agentspeak.grounded(term.args[1], intention.scope)
            goal = agentspeak.grounded(term.args[2], intention.scope)
            self.add_behavior(self._send_to_llm(str(message_text), str(norms), str(goal)))
            yield

        @self.actions.add(".notify_norms", 1)
        def _notify_norms(agent, term, intention):
            norms = agentspeak.grounded(term.args[0], intention.scope)

            norm_update_message = InternalMessage(
                to=settings.agent_settings.user_interrupt_name,
                thread="active_norms_update",
                body=str(norms),
            )

            self.add_behavior(self.send(norm_update_message, should_log=False))
            yield

        @self.actions.add(".say", 1)
        def _say(agent, term, intention):
            """
            Make the robot say the given text instantly.
            """
            message_text = agentspeak.grounded(term.args[0], intention.scope)

            self.logger.debug('"say" action called with text=%s', message_text)

            speech_command = SpeechCommand(data=message_text)
            speech_message = InternalMessage(
                to=settings.agent_settings.robot_speech_name,
                sender=settings.agent_settings.bdi_core_name,
                body=speech_command.model_dump_json(),
            )

            self.add_behavior(self.send(speech_message))

            chat_history_message = InternalMessage(
                to=settings.agent_settings.llm_name,
                thread="assistant_message",
                body=str(message_text),
            )

            experiment_logger.chat(str(message_text), extra={"role": "assistant"})

            self.add_behavior(self.send(chat_history_message))

            yield

        @self.actions.add(".gesture", 2)
        def _gesture(agent, term, intention):
            """
            Make the robot perform the given gesture instantly.
            """
            gesture_type = agentspeak.grounded(term.args[0], intention.scope)
            gesture_name = agentspeak.grounded(term.args[1], intention.scope)

            self.logger.debug(
                '"gesture" action called with type=%s, name=%s',
                gesture_type,
                gesture_name,
            )

            if str(gesture_type) == "single":
                endpoint = RIEndpoint.GESTURE_SINGLE
            elif str(gesture_type) == "tag":
                endpoint = RIEndpoint.GESTURE_TAG
            else:
                self.logger.warning("Gesture type %s could not be resolved.", gesture_type)
                endpoint = RIEndpoint.GESTURE_SINGLE

            gesture_command = GestureCommand(endpoint=endpoint, data=gesture_name)
            gesture_message = InternalMessage(
                to=settings.agent_settings.robot_gesture_name,
                sender=settings.agent_settings.bdi_core_name,
                body=gesture_command.model_dump_json(),
            )
            self.add_behavior(self.send(gesture_message))
            yield

        @self.actions.add(".notify_user_said", 1)
        def _notify_user_said(agent, term, intention):
            user_said = agentspeak.grounded(term.args[0], intention.scope)

            msg = InternalMessage(
                to=settings.agent_settings.llm_name, thread="user_message", body=str(user_said)
            )

            self.add_behavior(self.send(msg))

            yield

        @self.actions.add(".notify_trigger_start", 1)
        def _notify_trigger_start(agent, term, intention):
            """
            Notify the UI about the trigger we just started doing.
            """
            trigger_name = agentspeak.grounded(term.args[0], intention.scope)

            self.logger.debug("Started trigger %s", trigger_name)

            msg = InternalMessage(
                to=settings.agent_settings.user_interrupt_name,
                sender=self.name,
                thread="trigger_start",
                body=str(trigger_name),
            )

            self.add_behavior(self.send(msg))

            yield

        @self.actions.add(".notify_trigger_end", 1)
        def _notify_trigger_end(agent, term, intention):
            """
            Notify the UI about the trigger we just started doing.
            """
            trigger_name = agentspeak.grounded(term.args[0], intention.scope)

            self.logger.debug("Finished trigger %s", trigger_name)

            msg = InternalMessage(
                to=settings.agent_settings.user_interrupt_name,
                sender=self.name,
                thread="trigger_end",
                body=str(trigger_name),
            )

            self.add_behavior(self.send(msg))

            yield

        @self.actions.add(".notify_goal_start", 1)
        def _notify_goal_start(agent, term, intention):
            """
            Notify the UI about the goal we just started chasing.
            """
            goal_name = agentspeak.grounded(term.args[0], intention.scope)

            self.logger.debug("Started chasing goal %s", goal_name)

            msg = InternalMessage(
                to=settings.agent_settings.user_interrupt_name,
                sender=self.name,
                thread="goal_start",
                body=str(goal_name),
            )

            self.add_behavior(self.send(msg))

            yield

        @self.actions.add(".notify_transition_phase", 2)
        def _notify_transition_phase(agent, term, intention):
            """
            Notify the BDI program manager about a phase transition.
            """
            old = agentspeak.grounded(term.args[0], intention.scope)
            new = agentspeak.grounded(term.args[1], intention.scope)

            msg = InternalMessage(
                to=settings.agent_settings.bdi_program_manager_name,
                thread="transition_phase",
                body=json.dumps({"old": str(old), "new": str(new)}),
            )

            self.add_behavior(self.send(msg))

            yield

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

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