import os
import requests

import streamlit as st

from langchain.callbacks.base import BaseCallbackHandler


@st.cache_resource(ttl="1h")
def upload_data(uploaded_files):
    files = {"file": uploaded_files}
    with st.spinner("Uploading PDF..."):
        response = requests.post(
            "http://127.0.0.1:8000/api/upload", files=files
        )

        if response.status_code == 200:
            st.success(
                f'{response.json()["message"][0]} Vector Store created successfully!'
            )
            st.session_state.uploaded_pdf = True
        else:
            st.error("Failed to upload PDF!")


class StreamHandler(BaseCallbackHandler):
    def __init__(
        self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""
    ):
        self.container = container
        self.text = initial_text
        self.run_id_ignore_token = None

    def on_llm_start(self, serialized: dict, prompts: list, **kwargs):
        # Workaround to prevent showing the rephrased question as output
        if prompts[0].startswith("Human"):
            self.run_id_ignore_token = kwargs.get("run_id")

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        if self.run_id_ignore_token == kwargs.get("run_id", False):
            return
        self.text += token
        self.container.markdown(self.text)


class PrintRetrievalHandler(BaseCallbackHandler):
    def __init__(self, container):
        self.status = container.status("**Context Retrieval**")

    def on_retriever_start(self, serialized: dict, query: str, **kwargs):
        self.status.write(f"**Question:** {query}")
        self.status.update(label=f"**Context Retrieval:** {query}")

    def on_retriever_end(self, documents, **kwargs):
        for idx, doc in enumerate(documents):
            source = os.path.basename(doc.metadata["source"])
            self.status.write(f"**Document {idx} from {source}**")
            self.status.markdown(doc.page_content)
        self.status.update(state="complete")