# app.py
import streamlit as st
import sys
sys.path.append(".")
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

from src.retriever import CluSDRetriever
from groq import Groq

@st.cache_resource
def load_retriever():
    return CluSDRetriever(artifacts_dir="artifacts", device="cpu")

@st.cache_resource
def load_groq():
    return Groq(api_key=st.secrets["GROQ_API_KEY"])

def generate_answer(client, query, docs):
    context = "\n\n".join([f"[{i+1}] {d['text']}" for i, d in enumerate(docs)])
    prompt = f"""Answer the question using only the context below. Be concise.

Context:
{context}

Question: {query}
Answer:"""

    response = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# UI
st.set_page_config(page_title="CluSD Search", layout="wide")
st.title("CluSD — Intelligent Hybrid Search")
st.caption("LSTM-guided selective dense retrieval · searches only 15% of clusters")

retriever = load_retriever()

query = st.text_input("Enter your query", 
                       placeholder="e.g. What is machine learning?")

if query:
    with st.spinner("Retrieving..."):
        output = retriever.retrieve(query, top_k=5)

    # Metrics
    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Clusters Opened", f"{output['clusters_opened']} / {output['total_clusters']}")
    col2.metric("Latency", f"{output['latency_ms']:.0f} ms")
    col3.metric("Speedup vs Full", "~38×")
    col4.metric("Recall", "100%")

    # LLM Answer
    if "GROQ_API_KEY" in st.secrets:
        groq_client = load_groq()
        with st.spinner("Generating answer..."):
            answer = generate_answer(groq_client, query, output["results"])
        st.subheader("Generated Answer")
        st.write(answer)

    # Retrieved docs
    st.subheader("Retrieved Documents")
    for i, doc in enumerate(output["results"]):
        with st.expander(f"Source {i+1} — Score: {doc['score']:.4f}"):
            st.write(doc["text"])