import anthropic
import streamlit as st
from streamlit.logger import get_logger
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.llms import OpenAI
from langchain.llms import HuggingFaceEndpoint
from langchain.chat_models import ChatAnthropic
from langchain.vectorstores import SupabaseVectorStore
from stats import add_usage

memory = ConversationBufferMemory(memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
openai_api_key = st.secrets.openai_api_key
anthropic_api_key = st.secrets.anthropic_api_key
hf_api_key = st.secrets.hf_api_key
logger = get_logger(__name__)

def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db, stats):
    
    if 'chat_history' not in st.session_state:
        st.session_state['chat_history'] = []
      
    query = st.text_area("## Ask a question (" + stats + " queries answered so far)", max_chars=500)
    columns = st.columns(2)
    with columns[0]:
        button = st.button("Ask")
    with columns[1]:
        clear_history = st.button("Clear History", type='secondary')
        
    st.markdown("---\n\n")
    
    if clear_history:
        # Clear memory in Langchain
        memory.clear()
        st.session_state['chat_history'] = []
        st.experimental_rerun()

    if button:
        qa = None
        add_usage(stats_db, "chat", "prompt" + query, {"model": model, "temperature": st.session_state['temperature']})
        if model.startswith("gpt"):
            logger.info('Using OpenAI model %s', model)
            qa = ConversationalRetrievalChain.from_llm(
                OpenAI(
                    model_name=st.session_state['model'], openai_api_key=openai_api_key, temperature=st.session_state['temperature'], max_tokens=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True)
        elif anthropic_api_key and model.startswith("claude"):
            logger.info('Using Anthropics model %s', model)
            qa = ConversationalRetrievalChain.from_llm(
                ChatAnthropic(
                    model=st.session_state['model'], anthropic_api_key=anthropic_api_key, temperature=st.session_state['temperature'], max_tokens_to_sample=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True, max_tokens_limit=102400)
        elif hf_api_key:
            logger.info('Using HF model %s', model)
            # print(st.session_state['max_tokens'])
            endpoint_url = ("https://api-inference.huggingface.co/models/"+ model)
            model_kwargs = {"temperature" : st.session_state['temperature'],
                            "max_new_tokens" : st.session_state['max_tokens'],
                            "return_full_text" : False}
            hf = HuggingFaceEndpoint(
                endpoint_url=endpoint_url,
                task="text-generation",
                huggingfacehub_api_token=hf_api_key,
                model_kwargs=model_kwargs
            )
            qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(search_kwargs={"score_threshold": 0.6, "k": 4,"filter": {"user": st.session_state["username"]}}), memory=memory, verbose=True, return_source_documents=True)
        
        print("Question>")
        print(query)
        st.session_state['chat_history'].append(("You", query))

        # Generate model's response and add it to chat history
        model_response = qa({"question": query})
        logger.info('Result: %s', model_response["answer"])
        sources = model_response["source_documents"]
        logger.info('Sources: %s', model_response["source_documents"])

        if len(sources) > 0:
            st.session_state['chat_history'].append(("Safety Copilot", model_response["answer"]))
        else:
            st.session_state['chat_history'].append(("Safety Copilot", "I am sorry, I do not have enough information to provide an answer. If there is a public source of data that you would like to add, please email copilot@securade.ai."))
        
        # Display chat history
        st.empty()
        chat_history = st.session_state['chat_history']
        for speaker, text in chat_history:
            st.markdown(f"**{speaker}:** {text}")