import streamlit as st
from langchain.callbacks.base import BaseCallbackHandler


class TokenCallback(BaseCallbackHandler):
    """
    This class implements a callback to write tokens to a chat message as they are produced by an
    AI agent.
    """

    def __init__(self):
        """
        Creates the callback handler.
        """
        super().__init__()

        # The message placeholder must be overriden by the calling app at each turn. It represents
        # an empty slot that will be filled with tokens produced by the agent on-the-fly.
        self._message_placeholder: st.container = None
        self._full_text = ""

    def reset(self, message_placeholder: st.container):
        """
        Clears previously saved tokens.

        :param message_placeholder: clean, empty container that will receive the new tokens.
        """

        self._message_placeholder = message_placeholder
        self._full_text = ""

    def on_llm_new_token(self, token: str, **kwargs):
        """
        Appends a new token to the current partially produced message.

        :param token: new token.
        :param kwargs: extra parameters from the superclass. Not used in this implementation.
        """

        self._full_text += token

        # Append a vertical bar in the end to create the impression that the agent is typing its
        # answer to the chat.
        if self._message_placeholder:
            self._message_placeholder.write(self._full_text + "|")

    # def on_chat_model_start(*args, **kwargs):
    #     """
    #     Empty method. This method is inherited by the super class but never implemented. No
    #     implementation is needed for our use cases. Copied here just to avoid warnings.
    #     """
    #     return super().on_chat_model_start(*args, **kwargs)
