import streamlit as st
from pandasai.llm.openai import OpenAI
from pandasai import PandasAI
from dotenv import load_dotenv
import os
import pandas as pd
import time
from datasets import load_dataset  # For Hugging Face datasets

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

# Initialize the LLM based on user selection
def initialize_llm(model_choice):
    if model_choice == "llama-3.3-70b":
        if not groq_api_key:
            st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
            return None
        return ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
    elif model_choice == "GPT-4o":
        if not openai_api_key:
            st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
            return None
        return OpenAI(api_token=openai_api_key)

# Load Hugging Face dataset
def load_huggingface_dataset(dataset_name):
    progress_bar = st.progress(0)
    try:
        progress_bar.progress(10)
        dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True, uniform_split=True)
        progress_bar.progress(50)
        if hasattr(dataset, "to_pandas"):
            df = dataset.to_pandas()
        else:
            df = pd.DataFrame(dataset)
        progress_bar.progress(100)
        return df
    except Exception as e:
        progress_bar.progress(0)
        raise e

# Load uploaded CSV file
def load_uploaded_csv(uploaded_file):
    progress_bar = st.progress(0)
    try:
        progress_bar.progress(10)
        time.sleep(1)
        progress_bar.progress(50)
        df = pd.read_csv(uploaded_file)
        progress_bar.progress(100)
        return df
    except Exception as e:
        progress_bar.progress(0)
        raise e

# Main function to handle interactions
def chat_with_csv(df, prompt, llm):
    pandas_ai = PandasAI(llm)
    result = pandas_ai.run(df, prompt=prompt)
    return result

# Streamlit app layout
st.set_page_config(layout="wide")
st.title("ChatCSV powered by LLM")

# Initialize session state for data storage
if "data" not in st.session_state:
    st.session_state.data = None

if "llm" not in st.session_state:
    st.session_state.llm = None

# Select LLM model
model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
st.session_state.llm = initialize_llm(model_choice)

if st.session_state.llm is not None:
    # Dataset selection
    input_option = st.radio(
        "Select Dataset Input:",
        ["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"],
        index=2,
        horizontal=True,
    )

    if input_option == "Use Repo Directory Dataset":
        file_path = "./source/test.csv"
        if st.button("Load Dataset"):
            try:
                with st.spinner("Loading dataset from the repo directory..."):
                    st.session_state.data = pd.read_csv(file_path)
                st.success(f"File loaded successfully from '{file_path}'!")
            except Exception as e:
                st.error(f"Error loading dataset from the repo directory: {e}")

    elif input_option == "Use Hugging Face Dataset":
        dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
        if st.button("Load Dataset"):
            try:
                st.session_state.data = load_huggingface_dataset(dataset_name)
                st.success(f"Hugging Face Dataset '{dataset_name}' loaded successfully!")
            except Exception as e:
                st.error(f"Error loading Hugging Face dataset: {e}")

    elif input_option == "Upload CSV File":
        input_csv = st.file_uploader("Upload a CSV File:", type=["csv"])
        if input_csv:
            try:
                st.session_state.data = load_uploaded_csv(input_csv)
                st.success("File uploaded successfully!")
            except Exception as e:
                st.error(f"Error reading uploaded file: {e}")

    # Display the loaded dataset
    if st.session_state.data is not None:
        st.dataframe(st.session_state.data, use_container_width=True)

        # Chat interface
        input_text = st.text_area("Enter your query")
        if input_text and st.button("Chat with CSV"):
            try:
                st.info("Your Query: " + input_text)
                result = chat_with_csv(st.session_state.data, input_text, st.session_state.llm)
                st.success(result)
            except Exception as e:
                st.error(f"Error processing your query: {e}")