# app.py
import streamlit as st
import numpy as np
import faiss
import os
from sentence_transformers import SentenceTransformer
import requests
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

st.set_page_config(page_title="Meraj Graphics Assistant")

# Load model and index
@st.cache_resource
def load_data():
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    index = faiss.read_index("index.faiss")
    chunks = np.load("chunks.npy", allow_pickle=True)
    return embedder, index, chunks

embedder, index, chunks = load_data()

# Search FAISS
def search(query, top_k=3):
    q_embed = embedder.encode([query])
    D, I = index.search(np.array(q_embed), top_k)
    results = [chunks[i] for i in I[0]]
    return "\n".join(results)

# Call Groq API
def query_groq(context, question):
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "llama3-8b-8192",
        "messages": [
            {"role": "system", "content": "Answer based on the context."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
        ]
    }
    response = requests.post(url, headers=headers, json=data)
    return response.json()["choices"][0]["message"]["content"]

# UI
st.title("📋 Meraj Graphics Chat Assistant")

question = st.text_input("Ask something about our services:")

if st.button("Get Answer"):
    if not question:
        st.warning("Please enter a question.")
    elif not GROQ_API_KEY:
        st.error("API key not found. Please set GROQ_API_KEY in your environment.")
    else:
        with st.spinner("Searching..."):
            context = search(question)
            answer = query_groq(context, question)
            st.success("Answer:")
            st.write(answer)