"""
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 logging
import re
import uuid
from collections.abc import AsyncGenerator

import httpx
from pydantic import ValidationError

from control_backend.agents import BaseAgent
from control_backend.core.agent_system import InternalMessage
from control_backend.core.config import settings

from ...schemas.llm_prompt_message import LLMPromptMessage
from .llm_instructions import LLMInstructions

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


class LLMAgent(BaseAgent):
    """
    LLM Agent.

    This agent is responsible for processing user text input and querying a locally
    hosted LLM for text generation. It acts as the conversational brain of the system.

    It receives :class:`~control_backend.schemas.llm_prompt_message.LLMPromptMessage`
    payloads from the BDI Core Agent, constructs a conversation history, queries the
    LLM via HTTP, and streams the response back to the BDI agent in natural chunks
    (e.g., sentence by sentence).

    :ivar history: A list of dictionaries representing the conversation history (Role/Content).
    """

    def __init__(self, name: str):
        super().__init__(name)
        self.history = []
        self._querying = False
        self._interrupted = False
        self._interrupted_message = ""
        self._go_ahead = asyncio.Event()

    async def setup(self):
        self.logger.info("Setting up %s.", self.name)

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

        Expects messages from :attr:`settings.agent_settings.bdi_core_name` containing
        an :class:`LLMPromptMessage` in the body.

        :param msg: The received internal message.
        """
        if msg.sender == settings.agent_settings.bdi_core_name:
            match msg.thread:
                case "prompt_message":
                    try:
                        prompt_message = LLMPromptMessage.model_validate_json(msg.body)
                        self.add_behavior(self._process_bdi_message(prompt_message))  # no block
                    except ValidationError:
                        self.logger.debug("Prompt message from BDI core is invalid.")
                case "assistant_message":
                    self._apply_conversation_message({"role": "assistant", "content": msg.body})
                case "user_message":
                    self._apply_conversation_message({"role": "user", "content": msg.body})
        elif msg.sender == settings.agent_settings.bdi_program_manager_name:
            if msg.body == "clear_history":
                self.logger.debug("Clearing conversation history.")
                self.history.clear()
        else:
            self.logger.debug("Message ignored.")

    async def _process_bdi_message(self, message: LLMPromptMessage):
        """
        Orchestrate the LLM query and response streaming.

        Iterates over the chunks yielded by :meth:`_query_llm` and forwards them
        individually to the BDI agent via :meth:`_send_reply`.

        :param message: The parsed prompt message containing text, norms, and goals.
        """
        if self._querying:
            self.logger.debug("Received another BDI prompt while processing previous message.")
            self._interrupted = True  # interrupt the previous processing
            await self._go_ahead.wait()  # wait until we get the go-ahead

        message.text = f"{self._interrupted_message} {message.text}"

        self._go_ahead.clear()
        self._querying = True
        full_message = ""
        async for chunk in self._query_llm(message.text, message.norms, message.goals):
            if self._interrupted:
                self._interrupted_message = message.text
                self.logger.debug("Interrupted processing of previous message.")
                break
            await self._send_reply(chunk)
            full_message += chunk
        else:
            self._querying = False

            self._apply_conversation_message(
                {
                    "role": "assistant",
                    "content": full_message,
                }
            )
            self.logger.debug(
                "Finished processing BDI message. Response sent in chunks to BDI core."
            )
            await self._send_full_reply(full_message)

        self._go_ahead.set()
        self._interrupted = False

    def _apply_conversation_message(self, message: dict[str, str]):
        if len(self.history) > 0 and message["role"] == self.history[-1]["role"]:
            self.history[-1]["content"] += " " + message["content"]
            return
        self.history.append(message)

    async def _send_reply(self, msg: str):
        """
        Sends a response message (chunk) back to the BDI Core Agent.

        :param msg: The text content of the chunk.
        """
        reply = InternalMessage(
            to=settings.agent_settings.bdi_core_name,
            sender=self.name,
            body=msg,
        )
        await self.send(reply)

    async def _send_full_reply(self, msg: str):
        """
        Sends a response message (full) to agents that need it.

        :param msg: The text content of the message.
        """
        message = InternalMessage(
            to=settings.agent_settings.text_belief_extractor_name,
            sender=self.name,
            body=msg,
        )
        await self.send(message)

    async def _query_llm(
        self, prompt: str, norms: list[str], goals: list[str]
    ) -> AsyncGenerator[str]:
        """
        Send a chat completion request to the local LLM service and stream the response.

        It constructs the full prompt using
            :class:`~control_backend.agents.llm.llm_instructions.LLMInstructions`.
        It streams the response from the LLM and buffers tokens until a natural break (punctuation)
        is reached, then yields the chunk. This ensures that the robot speaks in complete phrases
        rather than individual tokens.

        :param prompt: Input text prompt to pass to the LLM.
        :param norms: Norms the LLM should hold itself to.
        :param goals: Goals the LLM should achieve.
        :yield: Fragments of the LLM-generated content (e.g., sentences/phrases).
        """
        instructions = LLMInstructions(norms if norms else None, goals if goals else None)
        messages = [
            {
                "role": "developer",
                "content": instructions.build_developer_instruction(),
            },
            *self.history,
        ]

        message_id = str(uuid.uuid4())

        try:
            full_message = ""
            current_chunk = ""
            async for token in self._stream_query_llm(messages):
                if self._interrupted:
                    return

                full_message += token
                current_chunk += token

                experiment_logger.chat(
                    full_message,
                    extra={"role": "assistant", "reference": message_id, "partial": True},
                )

                # Stream the message in chunks separated by punctuation.
                # We include the delimiter in the emitted chunk for natural flow.
                pattern = re.compile(r".*?(?:,|;|:|—|–|\.{3}|…|\.|\?|!)\s*", re.DOTALL)
                for m in pattern.finditer(current_chunk):
                    chunk = m.group(0)
                    if chunk:
                        yield current_chunk
                        current_chunk = ""

            # Yield any remaining tail
            if current_chunk:
                yield current_chunk

            experiment_logger.chat(
                full_message,
                extra={"role": "assistant", "reference": message_id, "partial": False},
            )
        except httpx.HTTPError as err:
            self.logger.error("HTTP error.", exc_info=err)
            yield "LLM service unavailable."
        except Exception as err:
            self.logger.error("Unexpected error.", exc_info=err)
            yield "Error processing the request."

    async def _stream_query_llm(self, messages) -> AsyncGenerator[str]:
        """
        Perform the raw HTTP streaming request to the LLM API.

        :param messages: The list of message dictionaries (role/content).
        :yield: Raw text tokens (deltas) from the SSE stream.
        :raises httpx.HTTPError: If the API returns a non-200 status.
        """
        async with httpx.AsyncClient() as client:
            async with client.stream(
                "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": messages,
                    "temperature": settings.llm_settings.chat_temperature,
                    "stream": True,
                },
            ) as response:
                response.raise_for_status()

                async for line in response.aiter_lines():
                    if not line or not line.startswith("data: "):
                        continue

                    data = line[len("data: ") :]
                    if data.strip() == "[DONE]":
                        break

                    try:
                        event = json.loads(data)
                        delta = event.get("choices", [{}])[0].get("delta", {}).get("content")
                        if delta:
                            yield delta
                    except json.JSONDecodeError:
                        self.logger.error("Failed to parse LLM response: %s", data)
