from langchain.agents import create_csv_agent
from langchain.chat_models import ChatOpenAI
from langchain.agents.agent_types import AgentType
import os
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
import tempfile

def main():
    # Load environment variables from .env file (if you're using one)
    load_dotenv()

    # Load the OpenAI API key from the environment
    openai_api_key = os.getenv("OPENAI_API_KEY")

    # Check if the OpenAI API key is set
    if openai_api_key is None or openai_api_key == "Financial Agent":
        st.error("OpenAI API Key is not set in the environment. Please check your environment variables.")
        return
    else:
        st.success("Financial Agent Started ")

    # Upload the CSV file
    csv_file = st.file_uploader("Upload a CSV file", type="csv")

    # If a file is uploaded
    if csv_file is not None:
        try:
            # Read the CSV file to show a preview
            df = pd.read_csv(csv_file, encoding='ISO-8859-1')
            st.write("Preview of the CSV file:")
            st.dataframe(df.head())  # Display the first few rows of the DataFrame
        except Exception as e:
            st.error(f"Error reading the CSV file: {str(e)}")
            return

        # Save the uploaded file to a temporary location
        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp_file:
            tmp_file.write(csv_file.getbuffer())
            tmp_file_path = tmp_file.name

        # Create the agent
        agent = create_csv_agent(
            ChatOpenAI(temperature=0, model="gpt-4-turbo"),
            tmp_file_path,  # Pass the path of the saved CSV file
            verbose=True,
            agent_type=AgentType.OPENAI_FUNCTIONS,
        )

        # Ask the user for a question
        user_question = st.text_input("Ask a question:")

        # If a question is entered, process it with the agent
        if user_question and user_question.strip():
            with st.spinner(text="In progress..."):
                try:
                    response = agent.run(user_question)
                    st.write(response)
                except Exception as e:
                    st.error(f"Error while processing your request: {str(e)}")
        else:
            st.warning("Please enter a valid question.")

if __name__ == "__main__":
    main()