from __future__ import annotations
from typing import Any, Dict, Optional
from pydantic import Field
from memory.buffer import ConversationBufferMemory
from memory.zep.zep_chat_message_history import ZepChatMessageHistory


class ZepMemory(ConversationBufferMemory):
    """Persist your chain history to the Zep Memory Server.

    The number of messages returned by Zep and when the Zep server summarizes chat
    histories is configurable. See the Zep documentation for more details.

    Documentation: https://docs.getzep.com

    Example:
        .. code-block:: python

        memory = ZepMemory(
                    session_id=session_id,  # Identifies your user or a user's session
                    url=ZEP_API_URL,        # Your Zep server's URL
                    api_key=<your_api_key>, # Optional
                    memory_key="history",   # Ensure this matches the key used in
                                            # chain's prompt template
                    return_messages=True,   # Does your prompt template expect a string
                                            # or a list of Messages?
                )
        chain = LLMChain(memory=memory,...) # Configure your chain to use the ZepMemory
                                              instance


    Note:
        To persist metadata alongside your chat history, your will need to create a
    custom Chain class that overrides the `prep_outputs` method to include the metadata
    in the call to `self.memory.save_context`.


    About Zep
    =========
    Zep provides long-term conversation storage for LLM apps. The server stores,
    summarizes, embeds, indexes, and enriches conversational AI chat
    histories, and exposes them via simple, low-latency APIs.

    For server installation instructions and more, see:
    https://docs.getzep.com/deployment/quickstart/

    For more information on the zep-python package, see:
    https://github.com/getzep/zep-python

    """

    chat_memory: ZepChatMessageHistory
    human_name: str = Field(default="Human")
    ai_name: str = Field(default="AI")
    auto_save: bool = Field(default=True)

    def __init__(
        self,
        session_id: str,
        url: str = "http://localhost:8000",
        api_key: Optional[str] = None,
        output_key: Optional[str] = None,
        input_key: Optional[str] = None,
        return_messages: bool = False,
        human_prefix: str = "Human",
        ai_prefix: str = "AI",
        memory_key: str = "history",
    ):
        """Initialize ZepMemory.

        Args:
            session_id (str): Identifies your user or a user's session
            url (str, optional): Your Zep server's URL. Defaults to
                                 "http://localhost:8000".
            api_key (Optional[str], optional): Your Zep API key. Defaults to None.
            output_key (Optional[str], optional): The key to use for the output message.
                                              Defaults to None.
            input_key (Optional[str], optional): The key to use for the input message.
                                              Defaults to None.
            return_messages (bool, optional): Does your prompt template expect a string
                                              or a list of Messages? Defaults to False
                                              i.e. return a string.
            human_prefix (str, optional): The prefix to use for human messages.
                                          Defaults to "Human".
            ai_prefix (str, optional): The prefix to use for AI messages.
                                       Defaults to "AI".
            memory_key (str, optional): The key to use for the memory.
                                        Defaults to "history".
                                        Ensure that this matches the key used in
                                        chain's prompt template.
        """

        chat_message_history = ZepChatMessageHistory(
            session_id=session_id,
            url=url,
            api_key=api_key,
        )

        super().__init__(
            chat_memory=chat_message_history,
            output_key=output_key,
            input_key=input_key,
            return_messages=return_messages,
            human_prefix=human_prefix,
            ai_prefix=ai_prefix,
            memory_key=memory_key,
        )

    def save_human_message(self, content: str):
        self.chat_memory.add_user_message(
            content,
            metadata={
                "author": self.human_name,
            },
        )

    def save_ai_message(self, content: str):
        self.chat_memory.add_ai_message(
            content,
            metadata={
                "author": self.ai_name,
            },
        )

    def save_context(
        self,
        inputs: Dict[str, Any],
        outputs: Dict[str, str],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Save context from this conversation to buffer.

        Args:
            inputs (Dict[str, Any]): The inputs to the chain.
            outputs (Dict[str, str]): The outputs from the chain.
            metadata (Optional[Dict[str, Any]], optional): Any metadata to save with
                                                           the context. Defaults to None

        Returns:
            None
        """
        if self.auto_save is False:
            return

        input_str, output_str = self._get_input_output(inputs, outputs)
        self.chat_memory.add_user_message(
            input_str,
            metadata={
                "author": self.human_name,
            },
        )
        self.chat_memory.add_ai_message(
            output_str,
            metadata={
                "author": self.ai_name,
            },
        )
