import logging
from datetime import datetime
from typing import Any, Optional

from langchain.callbacks.base import BaseCallbackHandler
from langchain.chat_models.base import BaseChatModel
from langchain.schema import AIMessage, ChatMessage, HumanMessage, SystemMessage
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai.chat_models import ChatOpenAI

from synthetic_agents.common.config import MESSAGE_DATETIME_FORMAT, settings
from synthetic_agents.common.functions import safe_placeholder_mapping
from synthetic_agents.model.constants import (
    DEFAULT_LLM_TEMPERATURE,
    DEFAULT_LLM_TOP_P,
    DEFAULT_OPENAI_CHAT_MODEL_NAME,
)
from synthetic_agents.model.entity.chat_memory import ChatMemory
from synthetic_agents.model.entity.life_memory import LifeMemory
from synthetic_agents.model.entity.memory import Memory
from synthetic_agents.prompt.builder import PromptBuilder

logger = logging.getLogger()


class PromptContainer:
    """
    A wrapper class to store a prompt used in a call to an LLM. LangChain does not return the
    prompt. We have to capture it using a callback function.
    """

    def __init__(self, prompt: str = ""):
        """
        Creates a prompt container.

        :param prompt: prompt to store in the container.
        """
        self.prompt = prompt


class LanguageModel:
    """
    This class represents a language model that receives a prompt and responds to it by deploying
    an LLM for language production.

    When developing a prompt template, the following placeholders will be populated with memories
    and last message the agent has to respond to. Not all the placeholders need to be present in
    the prompt template, but if they are present, they will be replaced with contents detailed
    below.

    {all_memories}: All memories passed to the language model, formatted as:
                    <creation timestamp> - <content>
    {life_memories}: All life memories passed to the language model, formatted as:
                    <creation timestamp> - <content>
    {chat_memories}: All chat memories passed to the language model, formatted as:
                    <creation timestamp> - <content>
    {last_message}: Last message in the chat the agent is responding to.
    {current_date}: A reference date informed to the language model representing the current date
                    and time the conversation is happening.


    The order of the memories is preserved.
    """

    def __init__(
        self,
        system_prompt_builder: PromptBuilder,
        agent_name: Optional[str] = None,
        llm_name: str = DEFAULT_OPENAI_CHAT_MODEL_NAME,
        temperature: float = DEFAULT_LLM_TEMPERATURE,
        top_p: float = DEFAULT_LLM_TOP_P,
        max_tokens: Optional[int] = None,
        callbacks: Optional[list[BaseCallbackHandler]] = None,
        api_key: str = settings.open_ai_api_key,
        text_streaming: bool = True,
        current_date: datetime = datetime.now(),
        chat_history_length: int = 0,
    ):
        """
        Creates a language model.

        :param system_prompt_builder: an object to build system prompts for the llm based on
            agent's attributes and a pre-defined template.
        :param agent_name: an optional agent name to add at the end of a prompt signaling the next
            message is produced by that agent.
        :param llm_name: name of the LLM model to use.
        :param temperature: temperature value for randomization of the responses.
        :param top_p: the top-p next tokens to consider for sampling.
        :param max_tokens: maximum number of tokens to be produced by the model. If undefined, it
            will be limited to the maximum capacity of the LLM.
        :param callbacks: callbacks to be passed to the LLM.
        :param api_key: key for authentication purposes when using the LLM API.
        :param text_streaming: whether to the LLM should generate the full text or as a stream of
            tokens. This needs to be set to True if one wants to capture tokens as they are
            produced with a callback.
        :param current_date: current date for reference. If it exists in the prompt template, the
            placeholder {current_date} will be filled with this date.
        :param chat_history_length: the number of previous chat messages to keep in memory and add
            to the prompt. It works as a secondary memory independent of the memory model.
        :raise ValueError: if max_tokens is 0 or a negative number.
        """

        if max_tokens and max_tokens <= 0:
            raise ValueError(
                f"The maximum number of tokens ({max_tokens}) must be a positive " f"number."
            )

        self.system_prompt_builder = system_prompt_builder
        self.agent_name = agent_name
        self.llm_name = llm_name
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.callbacks = callbacks
        self.api_key = api_key
        self.text_streaming = text_streaming
        self.current_date = current_date
        self.chat_history_length = chat_history_length

        # To be instantiated in a call to initialize_model
        self._llm: Optional[BaseChatModel] = None
        self._last_prompt = PromptContainer()

        # Memories and last message delivered to be used in the prompt to the LLM if defined in the
        # prompt template. Proper values can be updated via the setter methods during the
        # course of an application.
        self._life_memory_content = ""
        self._chat_memory_content = ""
        self._all_memory_content = ""
        self._last_message = ""

        self._chat_history: list[ChatMessage] = []

        if self.callbacks is None:
            self.callbacks = []

        # Custom callback so we can retrieve the prompt used by an LLM call.
        class CustomHandler(BaseCallbackHandler):
            def __init__(self, prompt_container):
                self.prompt_container = prompt_container

            def on_llm_start(
                self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any
            ) -> Any:
                self.prompt_container.prompt = "\n".join(prompts)

        self.callbacks.append(CustomHandler(self._last_prompt))

    def set_memories(self, memories: list[Memory]):
        """
        Stores life and chat memory contents to be passed to the internal LLM as strings.

        :param memories: memories to be parsed.
        """
        life_memories = []
        chat_memories = []
        all_memories = []

        for memory in memories:
            if isinstance(memory, LifeMemory):
                life_memories.append(str(memory))
            elif isinstance(memory, ChatMemory):
                chat_memories.append(str(memory))
            else:
                logger.warning(
                    f"Memory ({memory}) of type ({memory.named_type}) detected but not used by "
                    f"the language model."
                )

            # We also store all memories without making a distinction. The experimenter will select
            # which ones to use via placeholders in the prompt.
            all_memories.append(str(memory))

        self._life_memory_content = "\n".join(life_memories)
        self._chat_memory_content = "\n".join(chat_memories)
        self._all_memory_content = "\n".join(all_memories)

    def set_context(self, context: str):
        """
        Any relevant context to be passed to the LLM.

        :param context: textual context.
        """
        self._context = context

    def set_last_message(self, last_message: str):
        """
        Stores the last message delivered in the application to which the agent is responding to.

        :param last_message: last message in the application.
        """
        self._last_message = last_message

    def initialize_model(self):
        """
        Instantiates a language production model.
        """

        if self.llm_name.startswith("gpt"):
            self._llm = ChatOpenAI(
                openai_api_key=self.api_key,
                model=self.llm_name,
                temperature=self.temperature,
                streaming=self.text_streaming,
                callbacks=self.callbacks,
                max_tokens=self.max_tokens,
                model_kwargs=dict(top_p=self.top_p),
            )
        elif self.llm_name.startswith("claude"):
            self._llm = ChatAnthropic(
                api_key=self.api_key,
                model=self.llm_name,
                temperature=self.temperature,
                streaming=self.text_streaming,
                callbacks=self.callbacks,
                max_tokens=self.max_tokens,
                top_p=self.top_p,
            )
        else:
            raise ValueError(f"The language model ({self.llm_name}) is not supported.")

    def generate_text(self) -> dict[str, str]:
        """
        Prompt used as input to the LLM. This function expects all placeholders have been filled
        and the prompt contains all information necessary to generate the next message.

        :raise ValueError: if the language production model is undefined.
        :return: a dictionary containing the text produced by the LLM and the prompt used as
            input. Example, {"text": "this is the text produced.", "prompt": "this is the prompt."}
        """

        if not self._llm:
            raise ValueError(
                "Language production model is undefined. Please, make sure to call "
                "the initialize_model function to create such a model before calling the "
                "generate_text function."
            )

        system_message = self.build_system_message()
        user_message = self._last_message if self._last_message else ""

        message = self._llm.invoke(
            [SystemMessage(system_message), HumanMessage(user_message)]
        ).content

        self._chat_history.append(HumanMessage(user_message))
        self._chat_history.append(AIMessage(message))
        while len(self._chat_history) > self.chat_history_length:
            self._chat_history.pop(0)

        return {"text": message, "prompt": self._last_prompt.prompt}

    def build_system_message(self) -> str:
        """
        Builds a prompt to be passed to the LLM.

        :return: prompt.
        """
        prompt = self.system_prompt_builder.build_prompt()
        prompt = safe_placeholder_mapping(
            template=prompt,
            placeholder_values={
                "all_memories": self._all_memory_content,
                "life_memories": self._life_memory_content,
                "chat_memories": self._chat_memory_content,
                "current_date": self.current_date.strftime(MESSAGE_DATETIME_FORMAT),
                "chat_history": ChatPromptTemplate.from_messages(self._chat_history).format(),
            },
        )
        return prompt
