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

# Initialize FAISS and Model
VECTOR_DB_PATH = "vector_database.faiss"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# Initialize embedding model
embedding_model = SentenceTransformer(EMBEDDING_MODEL)

# Load FAISS Index
def load_faiss():
    if not os.path.exists(VECTOR_DB_PATH):
        st.error("Vector database not found! Ensure FAISS index is created.")
        return None
    index = faiss.read_index(VECTOR_DB_PATH)
    return index

faiss_index = load_faiss()

# GROQ API setup
GROQ_API_KEY = "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"  # Set in Hugging Face secrets
client = Groq(api_key=GROQ_API_KEY)

MODEL_ID = "deepseek-r1-distill-llama-70b"

# Function to get nearest neighbor from FAISS
def search_faiss(query, top_k=3):
    query_embedding = embedding_model.encode(query, convert_to_numpy=True).reshape(1, -1)
    distances, indices = faiss_index.search(query_embedding, top_k)
    return indices

# Function to call DeepSeek model from GROQ
def generate_response(context, query):
    prompt = f"Use the following retrieved context to answer the question:\n\nContext:\n{context}\n\nQuestion: {query}\nAnswer:"
    response = client.chat.completions.create(
        model=MODEL_ID,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.choices[0].message.content

# Streamlit UI
st.title("💡 AI Chat with FAISS & GROQ")
st.write("Ask a question and get responses based on stored knowledge!")

query = st.text_input("🔍 Enter your query:")
if query:
    if faiss_index is None:
        st.error("FAISS database not loaded. Please check deployment.")
    else:
        indices = search_faiss(query)
        retrieved_context = "\n".join([f"Chunk {i}: Retrieved text" for i in indices[0]])
        response = generate_response(retrieved_context, query)
        st.write("### 🤖 AI Response:")
        st.write(response)