import datetime
import os

import gradio as gr
import langchain
import pickle
from langchain.llms import HuggingFaceHub

from chain import get_new_chain1

api_token = os.environ["HF_TOKEN"]

def get_faiss_store():
    with open("docs.pkl", 'rb') as f:
        faiss_store = pickle.load(f)
        return faiss_store


def load_model():

    print(langchain.__file__)

    vectorstore = get_faiss_store()

    flan_ul = HuggingFaceHub(repo_id="google/flan-ul2", 
                                model_kwargs={"temperature":0.1, "max_new_tokens":200},
                                huggingfacehub_api_token=api_token)

    qa_chain = get_new_chain1(vectorstore, flan_ul, flan_ul, isFlan=True)
    return qa_chain


def chat(inp, history, agent):
    history = history or []
    if agent is None:
        history.append((inp, "Please wait for model to load (3-5 seconds)"))
        return history, history
    print("\n==== date/time: " + str(datetime.datetime.now()) + " ====")
    print("inp: " + inp)
    history = history or []
    output = agent({"question": inp, "chat_history": history})
    answer = output["answer"]
    history.append((inp, answer))
    print(history)
    return history, history


block = gr.Blocks(css=".gradio-container {background-color: lightgray}")

with block:
    with gr.Row():
        gr.Markdown("<h3><center>🤗 Transformers Chat</center></h3><p>Ask questions about the Hugging Face Transformers Library, Powered by Flan-UL2</p>")

    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="What's the answer to life, the universe, and everything?",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            "How do I install transformers?",
            "How do I load pretrained instances with an AutoClass?",
            "How do I fine-tune a pretrained model?",
        ],
        inputs=message,
    )

    gr.HTML(
        """
    This simple application uses Langchain, an open-source LLM, and FAISS to do Q&A over the 🤗 Transformers Documentation."""
    )

    gr.HTML(
        "<center>Powered by <a href='huggingface.co'>Hugging Face 🤗</a> and <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>"
    )

    state = gr.State()
    agent_state = gr.State()

    block.load(load_model, inputs=None, outputs=[agent_state])

    submit.click(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
    message.submit(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])

block.launch(debug=True)