import gradio as gr
import pickle
from chain import get_chain
from ingest import ingest_data

global vectorstore
vectorstore = ingest_data()

# with open("db.pkl", "rb") as f:
#     global vectorstore
#     vectorstore = pickle.load(f)

chain = get_chain()

def chat(inp, history):
    history = history or []
    output = chain({'input_documents': vectorstore.similarity_search(inp, k=4), 'query': inp}, return_only_outputs=True)
    response = output['output_text']
    history.append((inp, response))
    return history, history


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

with block:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("""<img src=https://huggingface.co/spaces/sooolee/beer-sommelier/resolve/main/images/octo.jpeg width=150></left>""")
        with gr.Column(scale=20):
            gr.Markdown("<h2><left>\"Hi! Welcome to Gourmet Haus Staudt!!\"</left></h2>")
            gr.Markdown("<h2><left>\"I'm your Beer Sommelier, and here to help you choose the right beer for you!\"</left></h2>")

    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="Do you have any recommendations for IPA with tropical fruite flavors?",
            lines=1,
        )
    with gr.Row():
        submit = gr.Button(value="Send", variant="secondary", size="lg")

    gr.Examples(
        examples=[
            "What does Bock taste like? Do you recommend it?",
            "I like Lagunitas IPA. Do you have anything similar?",
            "Do you have any recommendations for Pilsner?",
            "What kind of German beers do you have?"
        ],
        inputs=message,
    )

    state = gr.State()

    message.submit(chat, inputs=[message, state], outputs=[chatbot, state])
    message.submit(lambda x: gr.update(value=''), [], [message])

    submit.click(chat, inputs=[message, state], outputs=[chatbot, state])
    submit.click(lambda x: gr.update(value=''), [], [message])


block.launch(debug=True)