from typing import Dict, List, Optional

from langchain.vectorstores.base import VectorStore
from sqlalchemy.orm import Session

from synthetic_agents.callback.model_trace import ModelTraceCallback
from synthetic_agents.callback.token import TokenCallback
from synthetic_agents.common.config import settings
from synthetic_agents.database.config import get_db, get_vector_db
from synthetic_agents.database.entity.agent import Agent as DBAgent
from synthetic_agents.model.brain import Brain
from synthetic_agents.model.constants import (
    DEFAULT_CHAT_HISTORY_LENGTH,
    DEFAULT_CONTEXT_WINDOW_CAPACITY,
    DEFAULT_LLM_TEMPERATURE,
    DEFAULT_LLM_TOP_P,
    DEFAULT_MEMORY_RETRIEVAL_CAPACITY,
    DEFAULT_OPENAI_CHAT_MODEL_NAME,
    DEFAULT_WORKING_MEMORY_TOKEN_CAPACITY,
    MAX_LLM_TOKENS,
)
from synthetic_agents.model.entity.world import WorldState
from synthetic_agents.model.language import LanguageModel
from synthetic_agents.model.memory import LifeMemory, Memory, MemoryModel
from synthetic_agents.prompt.builder import PromptBuilder


class Agent:
    """
    This class represents an agent and it's brain which is comprised of three other sub-models:
    language and memory. The agent uses its brain to update its perception about the world
    and generate new messages.

    At the current stage, each agent's functions calls the corresponding functions in the brain,
    so, the brain and agent classes could be merged if needed, but we prefer to design them
    separately to give more flexibility to future developments in which the agent can have extra
    non-cognitive functionalities. Moreover, the composition design pattern here is beneficial if
    one is to experiment changing the agent's brain with another.
    """

    def __init__(
        self,
        agent_id: int,
        agent_type: str,
        application_type: str,
        agent_attributes: Dict[str, str],
        prompt_builder: PromptBuilder,
        memory_db: Session,
        memory_embedding_db: VectorStore,
        initial_world_state: WorldState,
        llm_name: str = DEFAULT_OPENAI_CHAT_MODEL_NAME,
        temperature: float = DEFAULT_LLM_TEMPERATURE,
        top_p: float = DEFAULT_LLM_TOP_P,
        chat_history_length: int = DEFAULT_CHAT_HISTORY_LENGTH,
        max_tokens: int = MAX_LLM_TOKENS,
        token_callback: TokenCallback = None,
        llm_logging_callback: ModelTraceCallback = None,
        memory_retrieval_capacity: int = DEFAULT_MEMORY_RETRIEVAL_CAPACITY,
        working_memory_token_capacity: int = DEFAULT_WORKING_MEMORY_TOKEN_CAPACITY,
        working_memory_buffer: Optional[List[Memory]] = None,
        remember_previous_sessions: bool = True,
        retrieve_chat_memories: bool = True,
        context_window_capacity: int = DEFAULT_CONTEXT_WINDOW_CAPACITY,
        api_key: str = settings.open_ai_api_key,
    ):
        """
        Creates an agent and its brain.

        :param agent_id: the ID of the agent.
        :param agent_type: the type of the agent.
        :param application_type: the type of the application.
        :param agent_attributes: attributes of an agent in json format. For instance, {'name':
            'Alex', 'demographics': 'age=20, gender=male', ...}.
        :param prompt_builder: an object that handles the construction of the prompts.
        :param memory_db: a relational database instance to store agents and memories.
        :param memory_embedding_db: a vector store instance used to store memory embeddings.
        :param initial_world_state: initial state of the world.
        :param llm_name: the name of the underlying LLM used by the agent's language model.
        :param temperature: a non-negative value that controls the randomness of the responses. If
            0, responses are deterministic. The higher the value, the more random the responses are
            which may lead to more hallucinations.
        :param top_p: the top-p next tokens to consider for sampling.
        :param chat_history_length: the length of the chat history.
        :param max_tokens: optional maximum number of tokens in the generated text. If undefined,
            the maximum number of tokens will be determined by the LLM specs.
        :param token_callback: optional callback function passed to the language model to keep
            track of tokens as they are produced.
        :param llm_logging_callback: optional callback function passed to the language model to
            keep track of logs produced by the internal LLM.
        :param memory_retrieval_capacity: the maximum number of memories to be read from the
            memory embedding database during memory retrieval.
        :param working_memory_token_capacity: the size of the working memory in number of tokens.
            We measure this in number of tokens because the working memory is used as part of the
            prompt used in the language model, and this is limited by the maximum number of tokens
            the language model can consume.
        :param working_memory_buffer: an optional list of memories to initialize the agent's
            working memory with.
        :param remember_previous_sessions: the agent can retrieve memories generated in any
            application session.
        :param retrieve_chat_memories: whether to retrieve chat memories or only life ones.
        :param context_window_capacity: the number of messages to use for context when retrieving
            memories. Defaults to DEFAULT_CONTEXT_WINDOW_CAPACITY.
        :param api_key: key for authentication purposes when using the LLM API.
        """

        # TODO: validate attributes against a schema.
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.application_type = application_type
        self.agent_attributes = agent_attributes

        # Save individual callback functions as local attributes for external access
        self.token_callback = token_callback
        self.llm_logging_callback = llm_logging_callback

        self.prompt_builder = prompt_builder
        self.memory_db = memory_db
        self.memory_embedding_db = memory_embedding_db
        self.initial_world_state = initial_world_state
        self.llm_name = llm_name
        self.temperature = temperature
        self.top_p = top_p
        self.chat_history_length = chat_history_length
        self.max_tokens = max_tokens
        self.token_callback = token_callback
        self.llm_logging_callback = llm_logging_callback
        self.memory_retrieval_capacity = memory_retrieval_capacity
        self.working_memory_token_capacity = working_memory_token_capacity
        self.working_memory_buffer = working_memory_buffer
        self.remember_previous_sessions = remember_previous_sessions
        self.retrieve_chat_memories = retrieve_chat_memories
        self.context_window_capacity = context_window_capacity
        self.api_key = api_key

        self.brain = None

    def initialize_agent(self):
        """
        Create and initialize the agent's brain.
        """

        callbacks = []
        if self.token_callback:
            callbacks.append(self.token_callback)
        if self.llm_logging_callback:
            callbacks.append(self.llm_logging_callback)

        language = LanguageModel(
            agent_name=self.name,
            system_prompt_builder=self.prompt_builder,
            llm_name=self.llm_name,
            temperature=self.temperature,
            top_p=self.top_p,
            max_tokens=self.max_tokens,
            callbacks=callbacks,
            api_key=self.api_key,
            text_streaming=False if self.token_callback is None else True,
            chat_history_length=self.chat_history_length,
        )

        memory = MemoryModel(
            agent_id=self.agent_id,
            retrieval_capacity=self.memory_retrieval_capacity,
            working_memory_token_capacity=self.working_memory_token_capacity,
            working_memory_buffer=self.working_memory_buffer,
            memory_db=self.memory_db,
            embedding_db=self.memory_embedding_db,
            llm_name=DEFAULT_OPENAI_CHAT_MODEL_NAME,  # For token count
            remember_previous_sessions=self.remember_previous_sessions,
            retrieve_chat_memories=self.retrieve_chat_memories,
        )

        self.brain = Brain(
            agent_name=self.name,
            language_model=language,
            memory_model=memory,
            context_window_capacity=self.context_window_capacity,
            initial_world_state=self.initial_world_state,
        )

        self.brain.initialize_brain()

    @property
    def working_memory_items(self) -> list[Memory]:
        """
        Gets list of items in the agent's working memory.

        :return: items in the agent's working memory.
        """
        return self.brain.memory_model.working_memory_buffer

    @property
    def name(self) -> str:
        """
        Gets agent's name from its attributes. If not present in the dictionary of agent's
        attributes, return the agent type and id.

        :return:
        """
        return self.agent_attributes.get("name", "Agent")

    def perceive_world(self, world_state: WorldState):
        """
        Invokes the agent's brain to perceive the world and update its internal models.

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

    def speak(self) -> Dict[str, str]:
        """
        Asks the brain 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.
        """
        return self.brain.speak()

    @staticmethod
    def create(
        application_type: str,
        agent_type: str,
        prompt_template_version: str,
        attributes: dict[str, str],
        life_facts: list[LifeMemory],
    ) -> int:
        """
        Creates an agent with a collection of attributes and life facts and persist to the
        database. We use a memory model object for persistence of the life fact memories such
        that they are also replicated in the embeddings database.

        :param application_type: type of the application the agent was designed for.
        :param agent_type: type of the agent.
        :param prompt_template_version: template version the agent uses.
        :param attributes: attributes of the agent.
        :param life_facts: live facts of the agent.
        :return: the ID of the newly created agent.
        """
        db = next(get_db())
        agent = DBAgent(
            application_type=application_type,
            agent_type=agent_type,
            prompt_template_version=prompt_template_version,
            attributes=attributes,
        )

        db.add(agent)
        db.flush()  # to update the agent's ID
        new_agent_id = agent.agent_id

        if len(life_facts) > 0:
            # Order memories by creation time before inserting them to the database to
            # keep things tidy (ids chronologically ordered).
            life_facts.sort(key=lambda m: m.creation_timestamp)
            memory_model = MemoryModel(
                agent_id=new_agent_id,
                memory_db=db,
                embedding_db=next(get_vector_db()),
            )

            memory_model.persist(memories=memory_model.fill_memory_ids(life_facts))
        else:
            db.commit()

        db.close()

        return new_agent_id
