import gradio as gr
import pandas as pd
import random
import time
from pandasai import PandasAI
from pandasai.llm.openai import OpenAI
import os

df = pd.DataFrame()
OPENAPI_KEY_VAL = ""
llm = ""
pandas_ai = ""

def process_file(file, OPENAPI_KEY):
    global df, OPENAPI_KEY_VAL, llm, pandas_ai
    df = pd.read_csv(file.name)
    OPENAPI_KEY_VAL = OPENAPI_KEY 
    
    llm = OpenAI(api_token = OPENAPI_KEY_VAL)
    pandas_ai = PandasAI(llm) 
    
    return df.head(2)

with gr.Blocks() as demo:
    
    with gr.Row():
        file = gr.inputs.File(type="file", label="Upload CSV")
        key = gr.Textbox(placeholder="Paste OPENAI API KEY")
        out = gr.DataFrame(type="pandas")
    
    btn = gr.Button("Execute")
    btn.click(fn=process_file, inputs=[file, key], outputs=out)
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        bot_message = "Hi"
        try:
            bot_message = pandas_ai.run(df, prompt=message)
        except Exception as e:
            print(e)
            bot_message = "I'm sorry, but as an AI language model, I am unable to provide a meaningful response to this question."
        # bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
        chat_history.append((message, bot_message))
        time.sleep(1)
        return "", chat_history

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

demo.launch()