import streamlit as st
from langchain.callbacks.streamlit.streamlit_callback_handler import (
    StreamlitCallbackHandler,
)
from langchain.schema.output import LLMResult
from sql_formatter.core import format_sql


class VectorSQLSearchDBCallBackHandler(StreamlitCallbackHandler):
    def __init__(self) -> None:
        self.progress_bar = st.progress(value=0.0, text="Writing SQL...")
        self.status_bar = st.empty()
        self.prog_value = 0
        self.prog_interval = 0.2

    def on_llm_start(self, serialized, prompts, **kwargs) -> None:
        pass

    def on_llm_end(
            self,
            response: LLMResult,
            *args,
            **kwargs,
    ):
        text = response.generations[0][0].text
        if text.replace(" ", "").upper().startswith("SELECT"):
            st.markdown("### Generated Vector Search SQL Statement \n"
                        "> This sql statement is generated by LLM \n\n")
            st.markdown(f"""```sql\n{format_sql(text, max_len=80)}\n```""")
            self.prog_value += self.prog_interval
            self.progress_bar.progress(
                value=self.prog_value, text="Searching in DB...")

    def on_chain_start(self, serialized, inputs, **kwargs) -> None:
        cid = ".".join(serialized["id"])
        self.prog_value += self.prog_interval
        self.progress_bar.progress(
            value=self.prog_value, text=f"Running Chain `{cid}`..."
        )

    def on_chain_end(self, outputs, **kwargs) -> None:
        pass


class VectorSQLSearchLLMCallBackHandler(VectorSQLSearchDBCallBackHandler):
    def __init__(self, table: str) -> None:
        self.progress_bar = st.progress(value=0.0, text="Writing SQL...")
        self.status_bar = st.empty()
        self.prog_value = 0
        self.prog_interval = 0.1
        self.table = table