"""
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 json

import httpx
from pydantic import BaseModel, ValidationError

from control_backend.agents.base import BaseAgent
from control_backend.agents.bdi.agentspeak_generator import AgentSpeakGenerator
from control_backend.core.agent_system import InternalMessage
from control_backend.core.config import settings
from control_backend.schemas.belief_list import BeliefList, GoalList
from control_backend.schemas.belief_message import Belief as InternalBelief
from control_backend.schemas.belief_message import BeliefMessage
from control_backend.schemas.chat_history import ChatHistory, ChatMessage
from control_backend.schemas.program import BaseGoal, SemanticBelief

type JSONLike = None | bool | int | float | str | list["JSONLike"] | dict[str, "JSONLike"]


class BeliefState(BaseModel):
    """
    Represents the state of inferred semantic beliefs.

    Maintains sets of beliefs that are currently considered true or false.
    """

    true: set[InternalBelief] = set()
    false: set[InternalBelief] = set()

    def difference(self, other: "BeliefState") -> "BeliefState":
        return BeliefState(
            true=self.true - other.true,
            false=self.false - other.false,
        )

    def union(self, other: "BeliefState") -> "BeliefState":
        return BeliefState(
            true=self.true | other.true,
            false=self.false | other.false,
        )

    def __sub__(self, other):
        return self.difference(other)

    def __or__(self, other):
        return self.union(other)

    def __bool__(self):
        return bool(self.true) or bool(self.false)


class TextBeliefExtractorAgent(BaseAgent):
    """
    Text Belief Extractor Agent.

    This agent is responsible for processing raw text (e.g., from speech transcription) and
    extracting semantic beliefs from it.

    It uses the available beliefs received from the program manager to try to extract beliefs from a
    user's message, sends and updated beliefs to the BDI core, and forms a ``user_said`` belief from
    the message itself.
    """

    def __init__(self, name: str):
        super().__init__(name)
        self._llm = self.LLM(self, settings.llm_settings.n_parallel)
        self.belief_inferrer = SemanticBeliefInferrer(self._llm)
        self.goal_inferrer = GoalAchievementInferrer(self._llm)
        self._current_beliefs = BeliefState()
        self._current_goal_completions: dict[str, bool] = {}
        self._force_completed_goals: set[BaseGoal] = set()
        self.conversation = ChatHistory(messages=[])

    async def setup(self):
        """
        Initialize the agent and its resources.
        """
        self.logger.info("Setting up %s.", self.name)

    async def handle_message(self, msg: InternalMessage):
        """
        Handle incoming messages. Expect messages from the Transcriber agent, LLM agent, and the
        Program manager agent.

        :param msg: The received message.
        """
        sender = msg.sender

        match sender:
            case settings.agent_settings.transcription_name:
                self.logger.debug("Received text from transcriber: %s", msg.body)
                self._apply_conversation_message(ChatMessage(role="user", content=msg.body))
                await self._user_said(msg.body)
                await self._infer_new_beliefs()
                await self._infer_goal_completions()
            case settings.agent_settings.llm_name:
                self.logger.debug("Received text from LLM: %s", msg.body)
                self._apply_conversation_message(ChatMessage(role="assistant", content=msg.body))
            case settings.agent_settings.bdi_program_manager_name:
                await self._handle_program_manager_message(msg)
            case _:
                self.logger.info("Discarding message from %s", sender)
                return

    def _apply_conversation_message(self, message: ChatMessage):
        """
        Save the chat message to our conversation history, taking into account the conversation
        length limit.

        :param message: The chat message to add to the conversation history.
        """
        length_limit = settings.behaviour_settings.conversation_history_length_limit
        self.conversation.messages = (self.conversation.messages + [message])[-length_limit:]

    async def _handle_program_manager_message(self, msg: InternalMessage):
        """
        Handle a message from the program manager: extract available beliefs and goals from it.

        :param msg: The received message from the program manager.
        """
        match msg.thread:
            case "beliefs":
                self._handle_beliefs_message(msg)
                await self._infer_new_beliefs()
            case "goals":
                self._handle_goals_message(msg)
                await self._infer_goal_completions()
            case "achieved_goals":
                self._handle_goal_achieved_message(msg)
            case "conversation_history":
                if msg.body == "reset":
                    self._reset_phase()
            case _:
                self.logger.warning("Received unexpected message from %s", msg.sender)

    def _reset_phase(self):
        """
        Delete all state about the current phase, such as what beliefs exist and which ones are
        true.
        """
        self.conversation = ChatHistory(messages=[])
        self.belief_inferrer.available_beliefs.clear()
        self._current_beliefs = BeliefState()
        self.goal_inferrer.goals.clear()
        self._current_goal_completions = {}

    def _handle_beliefs_message(self, msg: InternalMessage):
        """
        Handle the message from the Program Manager agent containing the beliefs that exist for this
        phase.
        :param msg: A list of beliefs.
        """
        try:
            belief_list = BeliefList.model_validate_json(msg.body)
        except ValidationError:
            self.logger.warning(
                "Received message from program manager but it is not a valid list of beliefs."
            )
            return

        available_beliefs = [b for b in belief_list.beliefs if isinstance(b, SemanticBelief)]
        self.belief_inferrer.available_beliefs = available_beliefs
        self.logger.debug(
            "Received %d semantic beliefs from the program manager: %s",
            len(available_beliefs),
            ", ".join(b.name for b in available_beliefs),
        )

    def _handle_goals_message(self, msg: InternalMessage):
        """
        Handle the message from the Program Manager agent containing the goals that exist for this
        phase.
        :param msg: A list of goals.
        """
        try:
            goals_list = GoalList.model_validate_json(msg.body)
        except ValidationError:
            self.logger.warning(
                "Received message from program manager but it is not a valid list of goals."
            )
            return

        # Use only goals that can fail, as the others are always assumed to be completed
        available_goals = {g for g in goals_list.goals if g.can_fail}
        available_goals -= self._force_completed_goals
        self.goal_inferrer.goals = available_goals
        self.logger.debug(
            "Received %d failable goals from the program manager: %s",
            len(available_goals),
            ", ".join(g.name for g in available_goals),
        )

    def _handle_goal_achieved_message(self, msg: InternalMessage):
        """
        Handle message that gets sent when goals are marked achieved from a user interrupt. This
        goal should then not be changed by this agent anymore.
        :param msg: List of goals that are marked achieved.
        """
        # NOTE: When goals can be marked unachieved, remember to re-add them to the goal_inferrer
        try:
            goals_list = GoalList.model_validate_json(msg.body)
        except ValidationError:
            self.logger.warning(
                "Received goal achieved message from the program manager, "
                "but it is not a valid list of goals."
            )
            return

        for goal in goals_list.goals:
            self._force_completed_goals.add(goal)
            self._current_goal_completions[f"achieved_{AgentSpeakGenerator.slugify(goal)}"] = True

        self.goal_inferrer.goals -= self._force_completed_goals

    async def _user_said(self, text: str):
        """
        Create a belief for the user's full speech.

        :param text: User's transcribed text.
        """
        belief_msg = InternalMessage(
            to=settings.agent_settings.bdi_core_name,
            sender=self.name,
            body=BeliefMessage(
                replace=[InternalBelief(name="user_said", arguments=[text])],
            ).model_dump_json(),
            thread="beliefs",
        )
        await self.send(belief_msg)

    async def _infer_new_beliefs(self):
        """
        Determine which beliefs hold and do not hold for the current conversation state. When
        beliefs change, a message is sent to the BDI core.
        """
        conversation_beliefs = await self.belief_inferrer.infer_from_conversation(self.conversation)

        new_beliefs = conversation_beliefs - self._current_beliefs
        if not new_beliefs:
            self.logger.debug("No new beliefs detected.")
            return

        self._current_beliefs |= new_beliefs

        belief_changes = BeliefMessage(
            create=list(new_beliefs.true),
            delete=list(new_beliefs.false),
        )

        message = InternalMessage(
            to=settings.agent_settings.bdi_core_name,
            sender=self.name,
            body=belief_changes.model_dump_json(),
            thread="beliefs",
        )
        await self.send(message)

    async def _infer_goal_completions(self):
        """
        Determine which goals have been achieved given the current conversation state. When
        a goal's achieved state changes, a message is sent to the BDI core.
        """
        goal_completions = await self.goal_inferrer.infer_from_conversation(self.conversation)

        new_achieved = [
            InternalBelief(name=goal, arguments=None)
            for goal, achieved in goal_completions.items()
            if achieved and self._current_goal_completions.get(goal) != achieved
        ]
        new_not_achieved = [
            InternalBelief(name=goal, arguments=None)
            for goal, achieved in goal_completions.items()
            if not achieved and self._current_goal_completions.get(goal) != achieved
        ]
        for goal, achieved in goal_completions.items():
            self._current_goal_completions[goal] = achieved

        if not new_achieved and not new_not_achieved:
            self.logger.debug("No goal achievement changes detected.")
            return

        belief_changes = BeliefMessage(
            create=new_achieved,
            delete=new_not_achieved,
        )
        message = InternalMessage(
            to=settings.agent_settings.bdi_core_name,
            sender=self.name,
            body=belief_changes.model_dump_json(),
            thread="beliefs",
        )
        await self.send(message)

    class LLM:
        """
        Class that handles sending structured generation requests to an LLM.
        """

        def __init__(self, agent: "TextBeliefExtractorAgent", n_parallel: int):
            self._agent = agent
            self._semaphore = asyncio.Semaphore(n_parallel)

        async def query(self, prompt: str, schema: dict, tries: int = 3) -> JSONLike | None:
            """
            Query the LLM with the given prompt and schema, return an instance of a dict conforming
            to this schema. Try ``tries`` times, or return None.

            :param prompt: Prompt to be queried.
            :param schema: Schema to be queried.
            :param tries: Number of times to try to query the LLM.
            :return: An instance of a dict conforming to this schema, or None if failed.
            """
            try_count = 0
            while try_count < tries:
                try_count += 1

                try:
                    return await self._query_llm(prompt, schema)
                except (httpx.HTTPError, json.JSONDecodeError, KeyError) as e:
                    if try_count < tries:
                        continue
                    self._agent.logger.exception(
                        "Failed to get LLM response after %d tries.",
                        try_count,
                        exc_info=e,
                    )

            return None

        async def _query_llm(self, prompt: str, schema: dict) -> JSONLike:
            """
            Query an LLM with the given prompt and schema, return an instance of a dict conforming
            to that schema.

            :param prompt: The prompt to be queried.
            :param schema: Schema to use during response.
            :return: A dict conforming to this schema.
            :raises httpx.HTTPStatusError: If the LLM server responded with an error.
            :raises json.JSONDecodeError: If the LLM response was not valid JSON. May happen if the
                response was cut off early due to length limitations.
            :raises KeyError: If the LLM server responded with no error, but the response was
                invalid.
            """
            async with self._semaphore:
                async with httpx.AsyncClient() as client:
                    response = await client.post(
                        settings.llm_settings.local_llm_url,
                        headers={"Authorization": f"Bearer {settings.llm_settings.api_key}"}
                        if settings.llm_settings.api_key
                        else {},
                        json={
                            "model": settings.llm_settings.local_llm_model,
                            "messages": [{"role": "user", "content": prompt}],
                            "response_format": {
                                "type": "json_schema",
                                "json_schema": {
                                    "name": "Beliefs",
                                    "strict": True,
                                    "schema": schema,
                                },
                            },
                            "reasoning_effort": "low",
                            "temperature": settings.llm_settings.code_temperature,
                            "stream": False,
                        },
                        timeout=30.0,
                    )
                    response.raise_for_status()

                    response_json = response.json()
                    json_message = response_json["choices"][0]["message"]["content"]
                    return json.loads(json_message)


class SemanticBeliefInferrer:
    """
    Infers semantic beliefs from conversation history using an LLM.
    """

    def __init__(
        self,
        llm: "TextBeliefExtractorAgent.LLM",
        available_beliefs: list[SemanticBelief] | None = None,
    ):
        self._llm = llm
        self.available_beliefs: list[SemanticBelief] = available_beliefs or []

    async def infer_from_conversation(self, conversation: ChatHistory) -> BeliefState:
        """
        Process conversation history to extract beliefs, semantically. The result is an object that
        describes all beliefs that hold or don't hold based on the full conversation.

        :param conversation: The conversation history to be processed.
        :return: An object that describes beliefs.
        """
        # Return instantly if there are no beliefs to infer
        if not self.available_beliefs:
            return BeliefState()

        n_parallel = max(1, min(settings.llm_settings.n_parallel - 1, len(self.available_beliefs)))
        all_beliefs: list[dict[str, bool | None] | None] = await asyncio.gather(
            *[
                self._infer_beliefs(conversation, beliefs)
                for beliefs in self._split_into_chunks(self.available_beliefs, n_parallel)
            ]
        )
        new_beliefs = BeliefState()
        # Collect beliefs from all parallel calls
        for beliefs in all_beliefs:
            if beliefs is None:
                continue
            # For each, convert them to InternalBeliefs
            for belief_name, belief_holds in beliefs.items():
                # Skip beliefs that were marked not possible to determine
                if belief_holds is None:
                    continue
                belief = InternalBelief(name=belief_name, arguments=None)
                if belief_holds:
                    new_beliefs.true.add(belief)
                else:
                    new_beliefs.false.add(belief)
        return new_beliefs

    @staticmethod
    def _split_into_chunks[T](items: list[T], n: int) -> list[list[T]]:
        """
        Split a list into ``n`` chunks, making each chunk approximately ``len(items) / n`` long.

        :param items: The list of items to split.
        :param n: The number of desired chunks.
        :return: A list of chunks each approximately ``len(items) / n`` long.
        """
        k, m = divmod(len(items), n)
        return [items[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n)]

    async def _infer_beliefs(
        self,
        conversation: ChatHistory,
        beliefs: list[SemanticBelief],
    ) -> dict[str, bool | None] | None:
        """
        Infer given beliefs based on the given conversation.
        :param conversation: The conversation to infer beliefs from.
        :param beliefs: The beliefs to infer.
        :return: A dict containing belief names and a boolean whether they hold, or None if the
            belief cannot be inferred based on the given conversation.
        """
        example = {
            "example_belief": True,
        }

        prompt = f"""{self._format_conversation(conversation)}

Given the above conversation, what beliefs can be inferred?
If there is no relevant information about a belief belief, give null.
In case messages conflict, prefer using the most recent messages for inference.

Choose from the following list of beliefs, formatted as `- <belief_name>: <description>`:
{self._format_beliefs(beliefs)}

Respond with a JSON similar to the following, but with the property names as given above:
{json.dumps(example, indent=2)}
"""

        schema = self._create_beliefs_schema(beliefs)

        return await self._llm.query(prompt, schema)

    @staticmethod
    def _create_belief_schema(belief: SemanticBelief) -> tuple[str, dict]:
        return AgentSpeakGenerator.slugify(belief), {
            "type": ["boolean", "null"],
            "description": belief.description,
        }

    @staticmethod
    def _create_beliefs_schema(beliefs: list[SemanticBelief]) -> dict:
        belief_schemas = [
            SemanticBeliefInferrer._create_belief_schema(belief) for belief in beliefs
        ]

        return {
            "type": "object",
            "properties": dict(belief_schemas),
            "required": [name for name, _ in belief_schemas],
        }

    @staticmethod
    def _format_message(message: ChatMessage):
        return f"{message.role.upper()}:\n{message.content}"

    @staticmethod
    def _format_conversation(conversation: ChatHistory):
        return "\n\n".join(
            [SemanticBeliefInferrer._format_message(message) for message in conversation.messages]
        )

    @staticmethod
    def _format_beliefs(beliefs: list[SemanticBelief]):
        return "\n".join(
            [f"- {AgentSpeakGenerator.slugify(belief)}: {belief.description}" for belief in beliefs]
        )


class GoalAchievementInferrer(SemanticBeliefInferrer):
    """
    Infers whether specific conversational goals have been achieved using an LLM.
    """

    def __init__(self, llm: TextBeliefExtractorAgent.LLM):
        super().__init__(llm)
        self.goals: set[BaseGoal] = set()

    async def infer_from_conversation(self, conversation: ChatHistory) -> dict[str, bool]:
        """
        Determine which goals have been achieved based on the given conversation.

        :param conversation: The conversation to infer goal completion from.
        :return: A mapping of goals and a boolean whether they have been achieved.
        """
        if not self.goals:
            return {}

        goals_achieved = await asyncio.gather(
            *[self._infer_goal(conversation, g) for g in self.goals]
        )
        return {
            f"achieved_{AgentSpeakGenerator.slugify(goal)}": achieved
            for goal, achieved in zip(self.goals, goals_achieved, strict=True)
        }

    async def _infer_goal(self, conversation: ChatHistory, goal: BaseGoal) -> bool:
        prompt = f"""{self._format_conversation(conversation)}

Given the above conversation, has the following goal been achieved?

Description of the goal: {goal.description or goal.name}

Answer with literally only `true` or `false` (without backticks)."""

        schema = {
            "type": "boolean",
        }

        return await self._llm.query(prompt, schema)
