import json
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.common.constants import AAI_QUESTIONS_ASSET_FILE
from synthetic_agents.common.functions import load_asset_file
from synthetic_agents.model.agent import Agent
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.memory import Memory
from synthetic_agents.prompt.builder import PromptBuilder


class AAIInterviewer(Agent):
    """
    This class represents an agent that conducts an AAI interview. This agent has no Brain and
    respective submodules. The produced text is deterministic and it comes from a preestablished
    set of questions. This agent will ask one question at a time.
    """

    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 AAI interviewer. The interviewer has a complete constructor interface for
        compatibility with the super class. However, it has no 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.
        """
        super().__init__(
            agent_id=agent_id,
            agent_type=agent_type,
            application_type=application_type,
            agent_attributes=agent_attributes,
            token_callback=token_callback,
            llm_logging_callback=llm_logging_callback,
            prompt_builder=prompt_builder,
            memory_db=memory_db,
            memory_embedding_db=memory_embedding_db,
            initial_world_state=initial_world_state,
            llm_name=llm_name,
            temperature=temperature,
            top_p=top_p,
            chat_history_length=chat_history_length,
            max_tokens=max_tokens,
            memory_retrieval_capacity=memory_retrieval_capacity,
            working_memory_token_capacity=working_memory_token_capacity,
            working_memory_buffer=working_memory_buffer,
            remember_previous_sessions=remember_previous_sessions,
            retrieve_chat_memories=retrieve_chat_memories,
            context_window_capacity=context_window_capacity,
            api_key=api_key,
        )

        self._aai_questions = []
        self._next_question_index = -1

    @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", "Interviewer")

    @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 []

    def initialize_agent(self):
        """
        Load AAI questions.
        """
        self._aai_questions = json.loads(load_asset_file(AAI_QUESTIONS_ASSET_FILE))["questions"]

    def perceive_world(self, world_state: WorldState):
        """
        Increments the index of the next question to be asked.

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

    def speak(self) -> Dict[str, str]:
        """
        Returns the next AAI question.

        :return: a dictionary containing the generated text (key = "text") and the prompt
            (key = "prompt") that was used to generate the text.
        """
        if self._next_question_index >= len(self._aai_questions):
            return None

        return {"prompt": "", "text": self._aai_questions[self._next_question_index]}
