import streamlit as st
import pandas as pd
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
from dotenv import load_dotenv
from datasets import load_dataset
import os

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

if not OPENAI_API_KEY:
    st.error("OpenAI API Key is missing. Make sure you set it in a .env file.")
    st.stop()

# Initialize OpenAI LLM
llm = OpenAI(api_token=OPENAI_API_KEY)

# App title and description
st.title("Patent Analytics: Chat With Your Dataset")
st.markdown(
    """
    Upload a CSV file or load a dataset from Hugging Face to:
    - Analyze data with natural language queries.
    - Visualize trends and insights (e.g., "Plot the number of patents filed per year").
    """
)

# Initialize session state for the dataframe
if "df" not in st.session_state:
    st.session_state.df = None

# Dataset input options
input_option = st.sidebar.radio(
    "Choose Dataset Input Method",
    options=["Use Hugging Face Dataset", "Upload CSV File"],
    index=0
)

# Dataset loading logic
if input_option == "Use Hugging Face Dataset":
    dataset_name = st.sidebar.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
    if st.sidebar.button("Load Dataset"):
        try:
            dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True)
            st.session_state.df = pd.DataFrame(dataset)
            st.sidebar.success(f"Dataset '{dataset_name}' loaded successfully!")
        except Exception as e:
            st.sidebar.error(f"Error loading dataset: {e}")
elif input_option == "Upload CSV File":
    uploaded_file = st.sidebar.file_uploader("Upload CSV File:", type=["csv"])
    if uploaded_file:
        try:
            st.session_state.df = pd.read_csv(uploaded_file)
            st.sidebar.success("File uploaded successfully!")
        except Exception as e:
            st.sidebar.error(f"Error loading file: {e}")

# Show the loaded dataframe preview
if st.session_state.df is not None:
    st.subheader("Dataset Preview")
    st.dataframe(st.session_state.df.head(10))

    # Create a SmartDataFrame for PandasAI
    chat_df = SmartDataframe(st.session_state.df, config={"llm": llm})

    # Input box for user questions
    question = st.text_input(
        "Ask a question about your data or request a visualization",
        placeholder="E.g., 'Which assignee has the most patents?' or 'Plot patent filings per year'",
    )

    if question:
        with st.spinner("Processing your request..."):
            try:
                # Chat with the dataframe
                response = chat_df.chat(question)

                # Detect visualizations in the query
                if "plot" in question.lower() or "graph" in question.lower():
                    st.write("### Visualization")
                else:
                    st.write("### Response")
                
                # Display response or plot
                st.write(response)
                st.success("Request processed successfully!")
            except Exception as e:
                st.error(f"An error occurred: {e}")
else:
    st.write("Upload a CSV file or load a dataset to get started.")