import streamlit as st
import pickle
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline

EMBEDDING_PATH = "quote_embeddings.pkl"

@st.cache_resource(show_spinner=False)
def load_data_and_models():
    # Load quotes and embeddings
    with open(EMBEDDING_PATH, "rb") as f:
        quotes, embeddings = pickle.load(f)

    # Initialize embedder
    embedder = SentenceTransformer('all-MiniLM-L6-v2')

    # Prepare FAISS index
    embeddings_np = embeddings.astype('float32')
    index = faiss.IndexFlatIP(embeddings_np.shape[1])
    faiss.normalize_L2(embeddings_np)
    index.add(embeddings_np)

    # Initialize text generation pipeline with smaller max length for speed
    generator = pipeline('text-generation', model='distilgpt2')

    return quotes, index, embedder, generator

quotes, index, embedder, generator = load_data_and_models()

def rag_search(query, top_k=3):
    q_emb = embedder.encode([query]).astype('float32')
    faiss.normalize_L2(q_emb)
    scores, indices = index.search(q_emb, top_k)
    context = "\n".join([f"{quotes[i]['quote']} — {quotes[i].get('author','Unknown')}" for i in indices[0]])
    prompt = f"Answer using these quotes:\n{context}\nQuestion: {query}\nAnswer:"
    outputs = generator(prompt, max_length=80, num_return_sequences=1, do_sample=False)
    answer = outputs[0]['generated_text'].split('Answer:')[-1].strip()
    return answer

st.title("🧠 RAG Quote-Based Q&A App")
user_query = st.text_input("💬 Ask something related to quotes:")

if user_query:
    with st.spinner("🔍 Searching quotes & generating response..."):
        answer = rag_search(user_query)
    st.markdown("### ✅ Answer:")
    st.write(answer)