from typing import Any, Dict, List

import streamlit as st
from langchain_core.outputs import LLMResult
from streamlit.external.langchain import StreamlitCallbackHandler


class ChatDataSelfQueryCallBack(StreamlitCallbackHandler):
    def __init__(self):
        super().__init__(st.container())
        self._current_thought = None
        self.progress_bar = st.progress(value=0.0, text="Executing ChatData SelfQuery CallBack...")

    def on_llm_start(
            self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        self.progress_bar.progress(value=0.35, text="Communicate with LLM...")
        pass

    def on_chain_end(self, outputs, **kwargs) -> None:
        if len(kwargs['tags']) == 0:
            self.progress_bar.progress(value=0.75, text="Searching in DB...")

    def on_chain_start(self, serialized, inputs, **kwargs) -> None:

        pass

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        st.markdown("### Generate filter by LLM \n"
                    "> Here we get `query_constructor` results \n\n")

        self.progress_bar.progress(value=0.5, text="Generate filter by LLM...")
        for item in response.generations:
            st.markdown(f"{item[0].text}")

        pass