import streamlit as st
import pandas as pd
from sqlalchemy import create_engine
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_openai import AzureChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
 
# Streamlit UI
def main():
    st.set_page_config(layout="wide")
    # st.title("Pharma POC")
    # Centered title
    st.markdown("""
        <h1 style='text-align: center;'>Pharma POC</h1>
    """, unsafe_allow_html=True)
 
    # Initialize session state for chat history
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "conversation_started" not in st.session_state:
        st.session_state.conversation_started = False
   
    # Sidebar for history and file upload
    with st.sidebar:
        # st.header("File Upload")
        uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
        if uploaded_file is not None:
            file_path = "temp_data.csv"
            with open(file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
           
            # Load CSV into DataFrame
            df = pd.read_csv(file_path)
           
            if not st.session_state.conversation_started:
                st.write("### Data Preview")
                st.dataframe(df.head())
           
            # Create SQL Database
            engine = create_engine("sqlite:///your_database.db")
            df.to_sql("your_table", engine, index=False, if_exists="replace")
            st.success("Data loaded into the database successfully!")
           
            # Load CSV using LangChain
            loader = CSVLoader(file_path=file_path)
            data = loader.load()
           
            # Initialize LLM
            api_key = "e88cfcf034ee4c53b5d0620504d69519"
            azure_endpoint = "https://transcendstore.openai.azure.com/"
            api_version = "2023-03-15-preview"
           
            llm = AzureChatOpenAI(
                model="gpt-4",
                deployment_name="gpt-4ogs",
                temperature=0,
                api_key=api_key,
                azure_endpoint=azure_endpoint,
                api_version=api_version,
            )
           
            db = SQLDatabase(engine=engine)
            agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True, top_k=50)
 
   
    # st.write("### Chat with Your Data")
    # Centered conversation headline
    # st.markdown("""
    #     <h3 style='text-align: center;'>Chat with Your Data</h3>
    # """, unsafe_allow_html=True)
    user_query = st.chat_input("Ask a question about your data:")
    if user_query:
        st.session_state.conversation_started = True
        response = agent_executor.invoke({"input": user_query})
        output_text = response.get("output", "No response available.")
 
        # Filter out table and column names for non-data-related queries
        if user_query.lower() in ["hi", "hello", "hey", "good morning", "good evening"]:
            output_text = "Hello! How can I assist you today?"
 
        st.session_state.chat_history.append((user_query, output_text))
       
    # Display chat history
    for query, response in st.session_state.chat_history:
        with st.chat_message("user"):
            st.write(query)
        with st.chat_message("assistant"):
            st.write(response)
 
if __name__ == "__main__":
    main()