import os
import json
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from groq import Groq
import gradio as gr

# Load Groq API key
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("❌ Missing GROQ_API_KEY. Please add it in your Hugging Face Space settings (Secrets).")

client = Groq(api_key=GROQ_API_KEY)

# Load FAISS index + metadata
DB_DIR = "kpi_vector_db"
index = faiss.read_index(os.path.join(DB_DIR, "kpi_index.faiss"))
with open(os.path.join(DB_DIR, "metadata.json"), "r", encoding="utf-8") as f:
    metadata = json.load(f)

# Load embedding model
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

def embed_text(text):
    return embed_model.encode([text])[0]

def retrieve(query, top_k=3):
    """Retrieve top_k chunks safely from FAISS + metadata."""
    q_emb = embed_text(query).astype("float32")
    D, I = index.search(np.array([q_emb]), top_k)

    results = []
    for idx in I[0]:
        idx = int(idx)  # ensure plain int
        if str(idx) in metadata:
            results.append(metadata[str(idx)])
        elif idx in metadata:
            results.append(metadata[idx])
    return results

def build_prompt(query, retrieved_chunks):
    context = "\n\n".join([chunk.get("text", "") for chunk in retrieved_chunks])
    system_prompt = "You are an AI assistant that answers questions based on company KPI Q3 documents (Excel and PPTX). Do not mention Q1 and Q2; mention Q3 if needed"

    user_message = f"Context:\n{context}\n\nQuestion: {query}\nAnswer in detail:"
    return system_prompt, user_message


def ask_groq(system_prompt, user_message):
    """Send the prompt to Groq LLaMA model."""
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",  # supported model
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    )
    return response.choices[0].message.content

def chatbot(query, history):
    """Main chatbot function for Gradio ChatInterface."""
    # Handle casual greetings without hitting FAISS
    if query.strip().lower() in ["hi", "hello", "hey"]:
        return "👋 Hello! I’m your KPI assistant. Ask me anything."

    retrieved = retrieve(query, top_k=3)
    if not retrieved:
        return "⚠️ Sorry, I couldn't find any relevant context in the documents."

    system_prompt, user_message = build_prompt(query, retrieved)
    answer = ask_groq(system_prompt, user_message)

    # Build safe sources list
    sources_list = []
    for c in retrieved:
        doc_name = c.get("doc", "Unknown document")
        chunk_id = c.get("chunk", "?")
        sources_list.append(f"- {doc_name} (chunk {chunk_id})")

    sources = "\n\nSources:\n" + "\n".join(sources_list)
    return answer

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📊 KPI Chatbot (Gradio + Groq)")
    chatbot_ui = gr.ChatInterface(fn=chatbot, type="messages")

if __name__ == "__main__":
    demo.launch()