import streamlit as st
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from google.colab import userdata # Import userdata

# ============================================
# Page Config
# ============================================
st.set_page_config(page_title="Flykite HR Bot", page_icon="✈️")

st.title("✈️ Flykite Airlines HR Assistant")
st.write("Ask any HR policy-related question")

# ============================================
# Load LLM
# ============================================
groq_api_key = userdata.get("GROQ_API_KEY") # Changed from st.secrets to userdata.get

llm = ChatGroq(
    model="openai/gpt-oss-120b",
    api_key=groq_api_key,
    temperature=0.3
)

# ============================================
# Load & Prepare Data (CACHE for speed)
# ============================================
@st.cache_resource
def load_vector_db():
    # Use the absolute path to the PDF file
    loader = PyPDFLoader("/content/drive/MyDrive/Dataset - Flykite Airlines_ HRP.pdf")
    documents = loader.load()

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100
    )
    docs = splitter.split_documents(documents)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vector_db = FAISS.from_documents(docs, embeddings)

    return vector_db

vector_db = load_vector_db()

retriever = vector_db.as_retriever(search_kwargs={"k": 3})

# ============================================
# RAG Function
# ============================================
def get_answer(question):
    docs = retriever.invoke(question)
    context = "\n\n".join([d.page_content for d in docs])

    prompt = f"""
You are an HR assistant for Flykite Airlines.

Answer ONLY from the context below.
If not found, say: "Information not found in policy."

Context:
{context}

Question:
{question}

Answer:
"""

    response = llm.invoke(prompt)
    return response.content, docs

# ============================================
# Chat UI
# ============================================
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

user_input = st.chat_input("Type your question here...")

if user_input:
    answer, docs = get_answer(user_input)

    st.session_state.chat_history.append(("user", user_input))
    st.session_state.chat_history.append(("bot", answer))

# ============================================
# Display Chat
# ============================================
for role, msg in st.session_state.chat_history:
    if role == "user":
        st.chat_message("user").write(msg)
    else:
        st.chat_message("assistant").write(msg)