import os
import pandas as pd
import streamlit as st
from pandasai import SmartDataframe, SmartDatalake
from pandasai.llm import GooglePalm
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_experimental.agents import create_pandas_dataframe_agent
# from pandasai.llm.openai import OpenAI
from langchain_groq import ChatGroq


def read_csv_file(uploaded_file):
    try:
        df = pd.read_csv(uploaded_file)
    except UnicodeDecodeError:
        st.warning("UTF-8 encoding failed. Trying latin1 encoding.")
        try:
            df = pd.read_csv(uploaded_file, encoding='latin1')
            print(uploaded_file)
        except UnicodeDecodeError:
            st.error("Unable to read the file with both UTF-8 and ISO-8859-1 encodings.")
            return None
    return df


def main():
    st.set_page_config(page_title="Smart Data Query App", layout="wide")

    st.title("Smart Data Query App")

    # Choice to upload one or two CSV files
    file_count = st.radio("How many CSV files would you like to upload?", (1, 2))

    col1, col2 = st.columns(2)
    with col1:
        uploaded_file1 = st.file_uploader("Choose the first CSV file", type="csv", key="file1")
    if file_count == 2:
        with col2:
            uploaded_file2 = st.file_uploader("Choose the second CSV file", type="csv", key="file2")
    else:
        uploaded_file2 = None

    if uploaded_file1 is not None or uploaded_file2 is not None:
        df1 = read_csv_file(uploaded_file1) if uploaded_file1 is not None else None
        df2 = read_csv_file(uploaded_file2) if uploaded_file2 is not None else None

        if df1 is not None or df2 is not None:
            if df1 is not None:
                with st.expander("Preview First CSV File"):
                    st.dataframe(df1.head())
            if df2 is not None:
                with st.expander("Preview Second CSV File"):
                    st.dataframe(df2.head())

            pandas_api = os.environ['PANDASAI_API_KEY']
            google_api = os.environ['GOOGLE_API_KEY']
            groq_api = os.environ['GROQ_API_KEY']

            # Set up the ChatGroq model
            llm = ChatGroq(
                groq_api_key=groq_api,
                model_name='mixtral-8x7b-32768'
            )
            # llm = GooglePalm(api_key=google_api)

            if df1 is not None and df2 is not None:
                lake = SmartDatalake([df1, df2])
            else:
                lake = SmartDataframe(df1, config = {"LLM": llm, "conversational": True, "verbose": True}) if df1 is not None else None

            datalake_1 = lake
            query = st.text_input("Enter your query:")

            submitted = st.button("Submit")

            if submitted:
                if query:
                    response = datalake_1.chat(query)

                    if "Unfortunately, I was not able to answer your question, because of the following error:" in response:
                        llm = ChatGoogleGenerativeAI(
                            model="gemini-pro", verbose=True, google_api_key=google_api
                        )

                        agent = create_pandas_dataframe_agent(
                            llm, 
                            df1,
                            verbose=True,
                        )

                        response2 = agent.invoke(query)
                        st.write("Response:")
                        st.write(response2["output"])

                    else:
                        st.write("Response:")
                        st.write(response)

                    # Append the query and response to the session state for query history
                    if 'query_history' not in st.session_state:
                        st.session_state.query_history = []
                    st.session_state.query_history.append((query, response))
                else:
                    st.write("Please enter a query.")

            # Display query history
            if 'query_history' in st.session_state and st.session_state.query_history:
                st.subheader("Query History")
                for q, r in st.session_state.query_history:
                    st.write(f"**Query:** {q}")
                    st.write(f"**Response:** {r}")
                    st.write("---")

        else:
            st.error("Failed to read one or both CSV files. Please check the files and try again.")
    else:
        st.write("Please upload at least one CSV file.")


if __name__ == "__main__":
    main()