from datetime import datetime
from typing import Dict

from synthetic_agents.model.entity.chat_memory import ChatMemory
from synthetic_agents.model.entity.world import WorldState
from synthetic_agents.model.language import LanguageModel
from synthetic_agents.model.memory import MemoryModel


class Brain:
    """
    This class represents a brain that comprises the following sub-models:
    - Language model
    - Memory model

    The brain is responsible for ensuring relevant information flows through the different
    sub-models for data persistence, world perception, emotion changes and language production.
    """

    def __init__(
        self,
        agent_name: str,
        language_model: LanguageModel,
        memory_model: MemoryModel,
        context_window_capacity: int,
        initial_world_state: WorldState,
    ):
        """
        Creates an agent's brain.

        :param agent_name: name of the agent for prepending messages with it.
        :param language_model: a model of language production.
        :param memory_model: a model of information persistence and retrieval.
        :param context_window_capacity: the number of messages to use for context when retrieving
            memories.
        :param initial_world_state: initial state of the world.
        """
        self.agent_name = agent_name
        self.language_model = language_model
        self.memory_model = memory_model
        self.context_window_capacity = context_window_capacity
        self.last_world_state = initial_world_state

        # This will accumulate world states to be used as context for the processing models.
        # For instance, the last N messages exchanged will be passed as input to the memory
        # Model for memory retrieval.
        self._context: list[WorldState] = []

    def initialize_brain(self):
        """
        Initialize internal models that need so.
        """
        self.language_model.initialize_model()

    def perceive_world(self, world_state: WorldState):
        """
        Updates sub-models with the most up-to-date state of the world.

        :raise Exception: if it could not read or write memories and their embeddings to the
            databases.
        :param world_state: state of the world.
        """
        self.last_world_state = world_state

        if world_state.last_text_message is None:
            # Currently the agent is limited to perceive text only. In the future, remove this
            # return if the agent can perceive other things.
            return

        self._context.append(world_state)
        if len(self._context) > self.context_window_capacity:
            self._context.pop(0)

        context_string = "/n".join(
            [s.last_text_message.content for s in self._context if s.last_text_message is not None]
        )

        self.memory_model.refresh_working_memory(
            context=context_string,
            application_id=world_state.application_id,
            session_id=world_state.session_id,
        )

        self.language_model.set_memories(self.memory_model.working_memory_buffer)

        # Set the last message prepended with the agent that delivered it. For instance,
        # "Peter: How are you today?
        self.language_model.set_last_message(world_state.last_text_message.content)

        self.language_model.current_date = world_state.current_date

        new_memories = [
            ChatMemory(
                memory_id=0,
                creation_timestamp=world_state.last_text_message.creation_timestamp,
                content=world_state.last_text_message.labeled_content,
                chat_id=world_state.application_id,
                session_id=world_state.session_id,
            )
        ]
        self.memory_model.persist(self.memory_model.fill_memory_ids(new_memories))

    def speak(self) -> Dict[str, str]:
        """
        Uses the language model to produce a message as a response to the current knowledge of the
        world.

        :return: a dictionary containing the generated text (key = "text") and the prompt
            (key = "prompt") that was used to generate the text.
        """
        # Add new generated text as a new memory as well
        new_message = self.language_model.generate_text()

        # Persist message as a memory
        new_memories = [
            ChatMemory(
                memory_id=0,
                creation_timestamp=datetime.now(),
                # preserve world state category
                content=f"{self.agent_name}: {new_message['text']}",
                chat_id=self.last_world_state.application_id,
                session_id=self.last_world_state.session_id,
            )
        ]
        self.memory_model.persist(self.memory_model.fill_memory_ids(new_memories))

        return new_message
