import gradio as gr
import os
from typing import List

from langchain.llms import OpenAIChat
from langchain.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
from langchain.chains import LLMChain, ChatVectorDBChain
from langchain.callbacks.base import CallbackManager
from langchain.chains.chat_vector_db.prompts import CONDENSE_QUESTION_PROMPT as SYNTHESIS_PROMPT
from langchain.document_loaders import DirectoryLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.faiss import FAISS


response_prompt_template = """You are an AI assistant for the machine learning monitoring startup Arthur.  You are 
given the following extracted parts of a long document and a question. If the question
includes a request for code, provide a code block directly from the documentation. Don't try to make up an answer. If the question is not about Arthur, politely inform them that
you are tuned to only answer questions about Arthur.
=========
Example 1:
Question: What data do I need to send to Arthur?
=========
**3. What if my data is proprietary? Can I still use Arthur?**
Yes! Arthur offers on-premises installation for customers with data security requirements. By integrating Arthur
into your business's on-premises stack, you can be confident that all security requirements are met while still
getting the benefits of the computation and analytics Arthur provides.
***
**4. What if I don’t have ground truth labels for my data? Or what if I will have the ground truth labels in the future, 
but they are not available yet?**
You don't need ground truth labels to log your model's inferences with Arthur. 
If your ground truth labels become available after your model's inferences, whether seconds later or years later, 
Arthur can link these new ground truth values to your model's past predictions, linking the new values by ID to 
their corresponding inferences already in the Arthur system.
In the meantime, Arthur’s data drift metrics can offer leading indicators of model underperformance to keep you 
covered if your ground truth labels are delayed or never become available.
***
=========
Answer in Markdown:
The data you need to get into Arthur is only the inference data - no ground truth is needed, since it can be uploaded
at a later time. Also, if you have proprietary data, you can install Arthur on-premises to keep your own data security protocols.
=========
Now the real question:
Question: {question}
=========
{context}
=========
Answer in Markdown:"""
RESPONSE_PROMPT = PromptTemplate(
    template=response_prompt_template, input_variables=["context", "question"]
)


def get_docs_vectorstore(dir_name):
    loader = DirectoryLoader(dir_name)
    raw_documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
    )
    documents = text_splitter.split_documents(raw_documents)
    embeddings = OpenAIEmbeddings()
    return FAISS.from_documents(documents, embeddings)


def get_langchain_agent(api_key):
    os.environ["OPENAI_API_KEY"] = api_key

    vectorstore = get_docs_vectorstore("files/arthur-docs-markdown")

    manager = CallbackManager([])
    question_manager = CallbackManager([])
    stream_manager = CallbackManager([])

    question_gen_llm = OpenAIChat(
        temperature=0,
        verbose=True,
        callback_manager=question_manager,
    )
    streaming_llm = OpenAIChat(
        streaming=True,
        callback_manager=stream_manager,
        verbose=True,
        temperature=0,
    )

    question_generator = LLMChain(
        llm=question_gen_llm, prompt=SYNTHESIS_PROMPT, callback_manager=manager
    )
    chat_response_generator = load_qa_chain(
        streaming_llm, chain_type="stuff", prompt=RESPONSE_PROMPT, callback_manager=manager
    )

    agent = ChatVectorDBChain(
        vectorstore=vectorstore,
        combine_docs_chain=chat_response_generator,
        question_generator=question_generator,
        callback_manager=manager,
        return_source_documents=True)

    os.environ["OPENAI_API_KEY"] = ""
    return agent


def get_source_doc(output):
    sources = output["source_documents"]
    assert len(sources) > 0
    source_document = sources[0]
    html_filename = source_document.metadata['source']
    source_doc_link = html_filename.replace('files/', '')
    source_doc_file = html_filename.replace('files/docs.arthur.ai/', '').replace('.html', '')
    with open(source_doc_file, 'r') as f:
        source_text = f.read()
    return source_text, source_doc_link


def chat(inp, history, agent):
    history = history or []
    result = agent({"question": inp, "chat_history": history})
    chat_result = result["answer"]
    source_doc, source_link = get_source_doc(result)
    response = ""
    for word in chat_result.split(" "):
        response += word + " "
        yield history + [(inp, response)], history + [(inp, response)], source_doc, source_link


def launch_ask_arthur(share=False):
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown("<h1><center>Ask Arthur</center></h1><br><h7><center>"
                        "This is an experimental document-retrieval question-answering system for asking questions about the Arthur product. <br> <br> Step 1: paste your OpenAI API key and hit the Register button."
                        "<br> <br> Step 2: type in the chat (shown on the left), and when you enter a message we fetch a relevant page from the Arthur documentation (shown on the right) "
                        "and prompt the Ask Arthur chatbot to answer your question using the page. "
                        "The chatbot's answers are entirely unverified, but it can sometimes offer a helpful summary of a "
                        "lot of information, or integrate information from multiple sources for you.</center></h7>")
            with gr.Column():
                openai_api_key_textbox = gr.Textbox(
                    placeholder="Paste your OpenAI API key (sk-...)",
                    show_label=False,
                    lines=1,
                    type="password",
                )
                submit_api_key_button = gr.Button(value="Register API Key", variant="secondary").style(full_width=False)
        with gr.Row().style():
            with gr.Column():
                chatbot = gr.Chatbot(
                    label="AskArthur chat history")
                message = gr.Textbox(
                    label="In the AskArthur chat, you can ask follow up questions or ask for clarifications!"
                          "\nReload the demo to change the topic of conversation and refresh the language model.",
                    placeholder="Enter your question here...",
                    lines=1,
                )
                submit_message = gr.Button(value="Send", variant="secondary").style(full_width=False)
                gr.Examples(label="Frequently asked questions about Arthur",
                    examples=[
                        "What default drift metrics does Arthur for deployed models?",
                        "How do I integrate Arthur with my existing ML stack?",
                        "How does Arthur monitor models without ground truth?",
                        "What if my data is proprietary, can I still use Arthur?",
                    ],
                    inputs=message,
                )
    
            with gr.Column():
                source_link = gr.Markdown()
                source_page = gr.Markdown()

        state = gr.State()
        agent_state = gr.State()
        submit_api_key_button.click(
            get_langchain_agent,
            inputs=[openai_api_key_textbox],
            outputs=[agent_state],
        )
        submit_message.click(chat, inputs=[message, state, agent_state], outputs=[chatbot, state, source_page, source_link])
        message.submit(chat, inputs=[message, state, agent_state], outputs=[chatbot, state, source_page, source_link])

    demo.queue().launch(share=share)


launch_ask_arthur()