import streamlit as st
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from transformers import pipeline
import time

# Set page config as the first Streamlit command
st.set_page_config(page_title="AI-Buddy Assistant", page_icon="🤖", layout="centered")

# Custom CSS
st.markdown("""
    <style>
        .stChat {
            font-size: 1.3rem;
        }
        .stTextArea textarea {
            font-size: 1.2rem;
        }
        .stMarkdown {
            font-size: 1.2rem;
        }
        .stButton button {
            background-color: #ff4b4b;
            color: white;
            font-size: 1.2rem;
            padding: 10px 20px;
            border-radius: 5px;
            border: none;
        }
        .stButton button:hover {
            background-color: #ff6565;
        }
    </style>
""", unsafe_allow_html=True)

# Cache the QA system initialization
@st.cache_resource
def initialize_qa_system():
    # Load PDF
    pdf_path = "ai_buddy.pdf"  # Update path accordingly
    data = PyPDFLoader(pdf_path).load()
    
    # Split documents
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=750, chunk_overlap=150)
    splits = text_splitter.split_documents(data)
    
    # Initialize embeddings
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )
    
    # Create vector store
    vector_db = Chroma.from_documents(
        documents=splits,
        embedding=embeddings,
        persist_directory="vector_db"
    )
    
    # Initialize QA pipeline
    qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
    
    return vector_db, qa_pipeline

# Initialize session state for vector database and QA pipeline
if 'vector_db' not in st.session_state:
    st.session_state.vector_db, st.session_state.qa_pipeline = initialize_qa_system()

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "Hello! How can I assist you today?"}
    ]

st.title("🤖 AI-Buddy Assistant")

# Sidebar configuration
st.sidebar.title("Want to know how AI helps in your profession?")
professions = ["Software Engineer", "Data Scientist", "Marketing Specialist", 
               "Financial Analyst", "Teacher", "Doctor", "Project Manager", 
               "Consultant", "Business Analyst", "Other"]
fields = ["IT", "Healthcare", "Education", "Finance", "Marketing", 
          "Engineering", "Sales", "Human Resources", "Consulting", "Other"]

profession = st.sidebar.selectbox("Choose Your Profession", professions)
field = st.sidebar.selectbox("Choose Your Field/Domain", fields)

if profession == "Other":
    profession = st.sidebar.text_input("Please specify your profession")
if field == "Other":
    field = st.sidebar.text_input("Please specify your field")

description = st.sidebar.text_area("About you", 
                                 placeholder="Briefly describe your role")

# Function to execute query with streaming response
def execute_query_with_stream(question):
    retriever = st.session_state.vector_db.as_retriever()
    retrieved_docs = retriever.get_relevant_documents(question)
    
    if not retrieved_docs:
        return "No relevant information found."
        
    context = " ".join([doc.page_content.strip() for doc in retrieved_docs])
    response = st.session_state.qa_pipeline(question=question, context=context)
    
    return response['answer']

# Handle sidebar submission
if st.sidebar.button("Get AI Insights"):
    prompt = f"""
    My profession is {profession} in the {field} field. 
    Here's about me: {description}. 
    How can AI and AI-Buddy help me in my role?
    """
    st.session_state.messages.append({"role": "user", "content": prompt})

# Chat interface
if prompt := st.chat_input("Ask me anything about AI-Buddy"):
    st.session_state.messages.append({"role": "user", "content": prompt})

# Display chat history and generate responses
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

if st.session_state.messages[-1]["role"] == "user":
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response = execute_query_with_stream(st.session_state.messages[-1]["content"])
            
            # Simulate streaming effect
            placeholder = st.empty()
            displayed_response = ""
            for char in response:
                displayed_response += char
                placeholder.write(displayed_response)
                time.sleep(0.02)
            
            st.session_state.messages.append(
                {"role": "assistant", "content": response}
            )