#!/usr/bin/env python3
"""
RAG Chatbot – Hugging Face Space
Carrega FAISS já existente (sem rebuild)
"""

import os
import gradio as gr
from openai import OpenAI

from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings


# =========================
# CONFIG
# =========================
INDEX_DIR = "vectorstore_faiss"
EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

TOP_K = 6
MAX_CONTEXT_CHARS = 4500

NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
NVIDIA_MODEL = "meta/llama-3.3-70b-instruct"

client = OpenAI(
    base_url=NVIDIA_BASE_URL,
    api_key=NVIDIA_API_KEY
) if NVIDIA_API_KEY else None


SYSTEM_PROMPT = """Você responde perguntas usando apenas o CONTEXTO recuperado.
Se não houver evidência suficiente, diga claramente.
Seja objetivo.
"""


# =========================
# LOAD FAISS EXISTENTE
# =========================
embedding = HuggingFaceEmbeddings(model_name=EMB_MODEL)

try:
    vectordb = FAISS.load_local(
        INDEX_DIR,
        embedding,
        allow_dangerous_deserialization=True
    )
    STATUS = "✅ Índice FAISS carregado com sucesso."
except Exception as e:
    vectordb = None
    STATUS = f"❌ Erro ao carregar índice: {str(e)}"


# =========================
# RAG FUNCTION
# =========================
def format_context(docs):
    context = "\n\n".join([d.page_content for d in docs])
    if len(context) > MAX_CONTEXT_CHARS:
        context = context[:MAX_CONTEXT_CHARS]
    return context


def chat(message, history):
    if not client:
        return "❌ Configure NVIDIA_API_KEY em Settings → Secrets."

    if not vectordb:
        return "❌ Índice FAISS não carregado."

    docs = vectordb.similarity_search(message, k=TOP_K)
    context = format_context(docs)

    completion = client.chat.completions.create(
        model=NVIDIA_MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"CONTEXTO:\n{context}\n\nPERGUNTA:\n{message}"}
        ],
        temperature=0.3,
        max_tokens=800,
    )

    return completion.choices[0].message.content


# =========================
# SUGGESTION CARDS
# =========================
SUGGESTIONS = [
    "Resuma os principais pontos do documento.",
    "Quais procedimentos são descritos?",
    "Liste requisitos ou obrigações mencionadas.",
    "Explique os termos técnicos utilizados.",
    "Há prazos ou datas importantes?",
    "Existe checklist operacional?",
    "Quais seções são mais relevantes?",
    "Há diferenças entre versões?"
]


# =========================
# UI
# =========================
with gr.Blocks(title="SOGETREL – RAG Assistant") as demo:

    gr.Markdown(f"""
    ## 📚 SOGETREL – Document Assistant

    {STATUS}
    """)

    gr.ChatInterface(
        fn=chat,
        examples=SUGGESTIONS,
        title="Assistant",
        description="Pergunte algo sobre os documentos."
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)