import streamlit as st
import pandas as pd
import os
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
from dotenv import load_dotenv
import tempfile
import matplotlib.pyplot as plt

# Load environment variables
load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

# Ensure OpenAI API key is provided
if not openai_api_key:
    st.error("OpenAI API key is not set. Please add it to a .env file.")
    st.stop()

# Initialize the LLM
llm = OpenAI(api_token=openai_api_key)

st.title("Chat with CSV File Using PandasAI")

uploaded_file = st.file_uploader("Upload a CSV file", type="csv")

if uploaded_file:
    # Read the CSV file
    df = pd.read_csv(uploaded_file)
    st.write("### Data Preview")
    st.dataframe(df.head(10))

    # Create SmartDataFrame
    chat_df = SmartDataframe(df, config={"llm": llm})

    st.write("### Chat with Your Data")
    user_query = st.text_input("Enter your question about the data:")

    if user_query:
        try:
            response = chat_df.chat(user_query)
            st.success(f"Response: {response}")
        except Exception as e:
            st.error(f"Error: {e}")

    st.write("### Generate and View Graphs")
    plot_query = st.text_input("Enter a query to generate a graph (e.g., 'Can you plot me a bar graph of total weekly sales for each store?'):")

    if plot_query:
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                # PandasAI can handle plotting
                chat_df.chat(plot_query)

                # Save and display the plot
                temp_plot_path = os.path.join(temp_dir, "plot.png")
                plt.savefig(temp_plot_path)
                st.image(temp_plot_path, caption="Generated Plot", use_column_width=True)

        except Exception as e:
            st.error(f"Error: {e}")

#st.write("### Instructions")
#st.markdown(
#    "1. Upload a CSV file to get started.\n"
#    "2. Enter a question to interact with the data.\n"
#    "3. Enter a query to generate and view graphs.\n"
#)