import gradio as gr
import pandas as pd
from text2csv import query_csv

df = pd.read_csv("sales_data.csv")
df['date'] = pd.to_datetime(df['date'], format='mixed')

def chat(message, history):
    response, fig, code = query_csv(df, message)

    if fig:
        return response, fig
    return response, None

with gr.Blocks() as demo:
    gr.Markdown("## CSV Chatbot")

    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox(placeholder="Ask about the sales data...")
            clear = gr.Button("Clear")

        with gr.Column():
            plot_output = gr.Image(label="Plot")

    def respond(message, chat_history):
        response, fig = chat(message, chat_history)
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": response})
        return "", chat_history, fig

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

demo.launch()