import gradio as gr
import time
import openai

from llama_index import StorageContext, load_index_from_storage

import pandas as pd

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

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

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>")
    gr.Markdown("<p align='center' style='font-size: 20px;'>Disclaimer: this is a fan-made project to highlight the work of Andrew Huberman. To support this project, please have a look at <a href='https://hubermanlab.com/'>Huberman Lab</a>.</p>")
    with gr.Row().style():
        with gr.Column(scale=1.0):
            openai_api_key = gr.Textbox(
                show_label=False,
                placeholder="Set your OpenAI API key here.",
                lines=1,
                type="password"
            ).style(container=False)
    with gr.Row().style():
        with gr.Column(scale=0.85):
            msg = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter.",
                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?",
                  "What are the benefits of walks after lunch?"
                ],
        inputs=msg
    )
    chatbot = gr.Chatbot().style(height=250)
    
    clear = gr.Button("Clear")

    def respond(openai_api_key, message, chat_history):
        if not openai_api_key:
            return "No OpenAI key provided, please provide one.", chat_history
        os.environ["OPENAI_API_KEY"] = openai_api_key
        index = load_index_from_storage(storage_context)
        query_engine = index.as_query_engine(similarity_top_k=3)
        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, [openai_api_key, msg, chatbot], [msg, chatbot])
    btn2.click(respond, [openai_api_key, msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()