import gradio as gr
import random
import time
import openai

from llama_index import StorageContext, load_index_from_storage
import os

import pandas as pd

openai.api_key = os.getenv("OPENAI_API_KEY")
passwd = os.getenv("PASSWD_SECRET")

df = pd.read_csv("original_huberman.csv")

storage_context = StorageContext.from_defaults(persist_dir="./storage")
# load index
index = load_index_from_storage(storage_context)

query_engine = index.as_query_engine(similarity_top_k=4)

def get_podcast_and_youtube(response):
    podcasts = []
    for node in response.source_nodes:
        podcast = node.node.extra_info["filename"].split("/")[-1].split(".")[0]
        podcasts.append(podcast)

    mask = df.podcast.apply(lambda x: x in podcasts)
    return df.loc[mask]

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>HuberChat</center></h1>")
    gr.Markdown("<p align='center'><img src='https://yt3.googleusercontent.com/5ONImZvpa9_hYK12Xek2E2JLzRc732DWsZMX2F-AZ1cTutTQLBuAmcEtFwrCgypqJncl5HrV2w=s900-c-k-c0x00ffffff-no-rj' height='50' width='95'></p>")
    gr.Markdown("<p align='center' style='font-size: 20px;'>Hi! I am Andrew HuberChat, a chatbot trained to answer neurobiology.</p>")
    with gr.Row().style():
        with gr.Column(scale=0.85):
            msg = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter. The url must contain the media type. e.g, https://example.com/example.jpg",
                lines=1,
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            btn2 = gr.Button("Send").style(full_height=True)
    gr.Examples(
        examples=["What is love?",
                  "Why should I get sunlight exposure?",
                  "How to be more motivated?"
                ],
        inputs=msg
    )
    chatbot = gr.Chatbot().style(height=250)
    
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        response = query_engine.query(message)
        bot_message = response.response
        for i, row in get_podcast_and_youtube(response).iterrows():
            bot_message += f"\n\n\n Source: {row.podcast} \n\n Link: {row.youtube_id}"
        chat_history.append((message, bot_message))
        time.sleep(1)
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    btn2.click(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(auth=("mithril", passwd))