# =========================================
# 1️⃣ Imports
# =========================================
import os
import pickle
import faiss
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer
from groq import Groq

# =========================================
# 2️⃣ Groq API
# =========================================
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
GROQ_MODEL = "llama-3.3-70b-versatile"

# =========================================
# 3️⃣ Load FAISS + chunks
# =========================================
index = faiss.read_index("faiss_index.bin")

with open("chunks.pkl", "rb") as f:
    all_chunks = pickle.load(f)

# Embedding model
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# =========================================
# 4️⃣ LLM Call Function (MULTILINGUAL)
# =========================================
def groq_query(prompt):
    completion = client.chat.completions.create(
        messages=[
            {
                "role": "system",
                "content": """
You are a strict university assistant.

RULES:
- Use ONLY the provided context
- If answer is not in context, say:
  "The information is not available in the provided documents. Please check https://uobs.edu.pk/"

LANGUAGE RULES:
- If user asks in English → reply in English
- If user asks in Urdu (اردو) → reply in Urdu
- If user writes in Roman Urdu (e.g. "department ka head kon hai") → reply in Roman Urdu
- Keep answer clear, correct and simple
"""
            },
            {"role": "user", "content": prompt}
        ],
        model=GROQ_MODEL,
        temperature=0,
    )
    return completion.choices[0].message.content

# =========================================
# 5️⃣ RAG FUNCTION
# =========================================
def rag_answer(query, k=3):

    if not query.strip():
        return "Please enter a valid question."

    # Embed query
    query_embedding = embedding_model.encode([query])
    query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
    query_embedding = query_embedding.astype("float32")

    # Search FAISS
    distances, indices = index.search(query_embedding, k)
    retrieved_texts = [all_chunks[i] for i in indices[0]]

    context = "\n\n".join(retrieved_texts)

    prompt = f"""
Use ONLY the context below to answer the question.

Context:
{context}

Question:
{query}
"""

    return groq_query(prompt)

# =========================================
# 6️⃣ Chat function
# =========================================
def respond(message, history):

    if not history:
        history = []

    answer = rag_answer(message)

    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": answer})

    return history, ""

# =========================================
# 7️⃣ Sample Questions
# =========================================
mock_questions = [
    "What is the focus of the Botany Department?",
    "Who is the Head of Chemistry Department?",
    "Computer Science department ka head kon hai?",
    "اردو اور انگلش میں جواب دیں: Education department ka goal kya hai?"
]

# =========================================
# 8️⃣ Gradio UI
# =========================================
with gr.Blocks() as demo:

    gr.Markdown(
        "<h1 style='text-align:center;'>🎓 UoBs Multilingual AI Assistant</h1>"
    )

    gr.Markdown("""
<div style="text-align:center;">
Ask questions in <b>English</b>, <b>Urdu</b>, or <b>Roman Urdu</b>.<br>
System will automatically respond in the same language.
</div>
    """)

    with gr.Row():

        # Left panel
        with gr.Column(scale=1):
            gr.Markdown("### 💡 Sample Questions")

            for q in mock_questions:
                btn = gr.Button(q)

                btn.click(
                    fn=lambda q=q: q,
                    outputs=None
                ).then(
                    lambda q=q: q,
                    outputs=None
                )

        # Right panel
        with gr.Column(scale=3):

            chatbot = gr.Chatbot(type="messages")

            with gr.Row():
                user_input = gr.Textbox(
                    placeholder="Ask in English, Urdu, or Roman Urdu...",
                    scale=8
                )

                send_btn = gr.Button("Send", scale=2)

            send_btn.click(
                respond,
                inputs=[user_input, chatbot],
                outputs=[chatbot, user_input]
            )

    demo.launch()