import streamlit as st
import pickle
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from groq import Groq
import os

# --- Load secrets ---
GROQ_API_KEY = st.secrets["GROQ_API_KEY"]

# --- Init Groq client ---
client = Groq(api_key=GROQ_API_KEY)

# --- Load embeddings model ---
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# --- Load docs + FAISS index ---
@st.cache_resource
def load_data():
    with open("documents.pkl", "rb") as f:
        documents = pickle.load(f)
    index = faiss.read_index("docs.index")
    return documents, index

documents, index = load_data()

# --- Retrieval ---
def retrieve(query, k=2):
    q_emb = embed_model.encode([query])
    D, I = index.search(np.array(q_emb), k)
    return [documents[i]["text"] for i in I[0]]

# --- RAG Query ---
def rag_query(query):
    retrieved = retrieve(query)
    context = "\n".join(retrieved)
    completion = client.chat.completions.create(
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
        ],
        model="llama-3.3-70b-versatile",
    )
    return completion.choices[0].message.content

# --- Streamlit UI ---
st.set_page_config(page_title="RAG with Groq", page_icon="📖")
st.title("📖 Simple RAG with Groq + SentenceTransformers")

query = st.text_input("Ask a question:")

if st.button("Submit") and query:
    with st.spinner("Thinking..."):
        answer = rag_query(query)
    st.subheader("Answer")
    st.write(answer)

    st.subheader("Retrieved context")
    for c in retrieve(query):
        st.markdown(f"- {c}")