app.py

# ---------- app.py ----------
# Dependências:
# pip install gradio faiss-cpu sentence-transformers openai

import os
from pathlib import Path
import pickle
from typing import List, Dict, Any

import gradio as gr
import faiss
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# ========= NVIDIA API =========
NV_API_KEY = os.environ.get("NVIDIA_API_KEY") or os.environ.get("NV_API_KEY")
if not NV_API_KEY:
    raise RuntimeError(
        "A chave da NVIDIA não foi encontrada.\n"
        "Defina um secret chamado 'NVIDIA_API_KEY' (ou NV_API_KEY) com a tua chave da NVIDIA.\n"
        "• Localmente: export NVIDIA_API_KEY='SUA_CHAVE'\n"
        "• Hugging Face Spaces: Settings -> Repository secrets -> Add secret."
    )

client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=NV_API_KEY,
)
CHAT_MODEL = "meta/llama3-8b-instruct"

# ========= Configuração do App =========
APP_TITLE = "EcoLexIA – Assistente Inteligente de Leis Ambientais de Portugal"

INTRO = (
    "👋 Bem-vindo ao **EcoLexIA**, o teu assistente jurídico especializado em **direito do ambiente em Portugal**.\n\n"
    "Este sistema utiliza **RAG (Retrieval-Augmented Generation)** para consultar automaticamente os documentos legais "
    "carregados (leis, decretos, regulamentos, pareceres, etc.) e responder às tuas perguntas com base nesses textos."
)

SUGGESTION_QUESTIONS = [
    "Resuma os principais princípios da Lei de Bases do Ambiente.",
    "Quais são as obrigações do Estado em matéria de proteção ambiental?",
    "Explique como funciona a Avaliação de Impacte Ambiental em Portugal.",
    "Que legislação regula a gestão de resíduos urbanos?",
    "Existe enquadramento legal para participação pública em decisões ambientais?",
    "Quais são as regras sobre emissões poluentes na indústria?",
]

SUGGESTIONS_THEMES = {
    "Lei de Bases do Ambiente": [
        "Quais são os princípios fundamentais da Lei de Bases do Ambiente?",
        "Como a Lei de Bases do Ambiente enquadra o desenvolvimento sustentável?",
    ],
    "Avaliação de Impacte Ambiental (AIA)": [
        "Explique o que é Avaliação de Impacte Ambiental e quando é obrigatória.",
        "Que entidades estão envolvidas no processo de AIA?",
    ],
    "Resíduos & Poluição": [
        "Que legislação trata da gestão de resíduos em Portugal?",
        "Que obrigações têm as empresas relativamente ao controlo de emissões poluentes?",
    ],
    "Ordenamento do Território & Conservação": [
        "Como o ordenamento do território se articula com a proteção ambiental?",
        "Que diplomas legais regulam áreas protegidas e conservação da natureza?",
    ],
}

# ========= Caminhos do índice =========
INDEX_FILE = "faiss_index.faiss"
EMBEDDINGS_FILE = "embeddings.pkl"

if not Path(INDEX_FILE).exists() or not Path(EMBEDDINGS_FILE).exists():
    raise FileNotFoundError(
        "❌ Índice não encontrado.\n"
        "Certifique-se de que 'faiss_index.faiss' e 'embeddings.pkl' "
        "foram gerados pelo build_index.py na mesma pasta deste app."
    )

index = faiss.read_index(INDEX_FILE)
with open(EMBEDDINGS_FILE, "rb") as f:
    emb_data = pickle.load(f)

texts = emb_data["texts"]
metadatas = emb_data["metadatas"]

# Mesmo modelo de embeddings usado no build_index.py
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# ✅ Gradio atual espera "messages format":
# lista de dicts: {"role": "user"/"assistant", "content": "..."}
dialog_history: List[Dict[str, str]] = []


# ========= Recuperação de contexto =========
def retrieve_context(query: str, k: int = 4) -> str:
    if not query or not query.strip():
        return ""

    q_emb = embedding_model.encode([query], convert_to_numpy=True)
    _, indices = index.search(q_emb, k)

    parts = []
    for idx in indices[0]:
        if idx < 0 or idx >= len(texts):
            continue
        chunk = texts[idx]
        meta = metadatas[idx] if idx < len(metadatas) else {}
        src = meta.get("source", "documento desconhecido")
        parts.append(f"[Documento: {src}]\n{chunk}")

    return "\n\n---\n\n".join(parts)


# ========= Streaming da NVIDIA =========
def nv_stream(messages: List[Dict[str, str]], temperature: float, top_p: float, max_tokens: int):
    reply = ""
    stream = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=messages,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
        stream=True,
    )
    for chunk in stream:
        delta = chunk.choices[0].delta
        if getattr(delta, "content", None):
            reply += delta.content
            yield reply


# ========= Lógica do chat =========
def chatbot(user_input: str, temperature: float, top_p: float, max_tokens: int):
    global dialog_history

    if not user_input or not user_input.strip():
        return dialog_history, ""

    context = retrieve_context(user_input, k=6)

    system_msg = {
        "role": "system",
        "content": (
            "És um assistente jurídico especializado em direito do ambiente em Portugal. "
            "Responde SEMPRE em português europeu, de forma clara e estruturada.\n\n"
            "Regras:\n"
            "1. Usa apenas o contexto abaixo para responder.\n"
            "2. Se não houver informação suficiente, diz que não encontras base nos documentos e "
            "sugere consultar a legislação oficial.\n"
            "3. Indica o nome do documento (PDF) sempre que fizer sentido.\n\n"
            f"=== CONTEXTO RECUPERADO ===\n{context}\n\n"
        ),
    }

    # Mensagens que vão para o modelo = system + histórico + user atual
    messages: List[Dict[str, str]] = [system_msg] + dialog_history + [{"role": "user", "content": user_input}]

    reply_full = ""
    try:
        for partial in nv_stream(messages, temperature, top_p, max_tokens):
            reply_full = partial

        # Atualiza histórico no formato messages (compatível com Gradio)
        dialog_history = dialog_history + [
            {"role": "user", "content": user_input},
            {"role": "assistant", "content": reply_full},
        ]
    except Exception as e:
        reply_full = f"⚠️ Erro na API NVIDIA: {type(e).__name__}: {e}"
        dialog_history = dialog_history + [
            {"role": "user", "content": user_input},
            {"role": "assistant", "content": reply_full},
        ]

    return dialog_history, ""


def clear_history():
    global dialog_history
    dialog_history = []
    return [], ""


# ========= CSS =========
custom_css = r"""
body, .gradio-container {
    background: #ffffff;
    font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}

/* Header azul */
#header-box {
    max-width: 1100px;
    margin: 1.5rem auto 1rem auto;
}

.header-card {
    background: linear-gradient(135deg, #0b3c91 0%, #1565c0 40%, #1e88e5 100%);
    border-radius: 16px;
    padding: 1.4rem 1.8rem;
    color: #ffffff;
    box-shadow: 0 14px 30px rgba(15, 23, 42, 0.18);
}

.header-title {
    font-size: 1.6rem;
    font-weight: 700;
    margin: 0;
    color: #ffffff !important;
}

.header-subtitle {
    margin: 0.35rem 0 0 0;
    font-size: 0.96rem;
    opacity: 0.95;
    color: #ffffff !important;
}

/* Cartões principais */
.card {
    background: #ffffff;
    border-radius: 16px;
    border: 1px solid #e0e0e0;
    box-shadow: 0 8px 18px rgba(15, 23, 42, 0.06);
    padding: 1rem 1.1rem;
}

/* Chat */
#chat-window {
    height: 60vh;
}

/* Botões */
.gr-button-primary {
    background: #1565c0 !important;
    color: #ffffff !important;
    border: none !important;
}

.gr-button-secondary {
    background: #f5f5f5 !important;
    color: #333333 !important;
    border: 1px solid #e0e0e0 !important;
}

/* Sugestões */
.suggestion-btn {
    width: 100%;
    justify-content: flex-start;
    font-size: 0.88rem;
}

/* Rodapé */
.app-footer {
    margin-top: 1rem;
    font-size: 0.8rem;
    text-align: center;
    color: #555555;
}
"""


# ========= Layout Gradio =========
with gr.Blocks(title=APP_TITLE) as demo:
    with gr.Group(elem_id="header-box"):
        gr.HTML(
            f"""
            <div class="header-card">
                <div class="header-title">{APP_TITLE}</div>
                <div class="header-subtitle">
                    Consultor jurídico inteligente com RAG sobre legislação ambiental portuguesa.
                </div>
            </div>
            """
        )

    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column(scale=3):
            with gr.Group(elem_classes="card"):
                gr.Markdown("### 💬 Conversa Jurídica")

                # ✅ Agora o valor/retorno é messages-format (dicts role/content)
                chatbot_ui = gr.Chatbot(
                    elem_id="chat-window",
                    label="Chatbot",
                )

                txt = gr.Textbox(
                    placeholder="Escreve aqui a tua pergunta sobre leis do ambiente em Portugal…",
                    lines=3,
                    show_label=False,
                )

                with gr.Row():
                    btn_send = gr.Button("Enviar", variant="primary")
                    btn_clear = gr.Button("Limpar", variant="secondary")

                with gr.Accordion("Parâmetros avançados", open=False):
                    temperature = gr.Slider(0, 1, value=0.5, label="Temperature")
                    top_p = gr.Slider(0, 1, value=0.9, label="Top-p")
                    max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max Tokens")

            btn_send.click(chatbot, [txt, temperature, top_p, max_tokens], [chatbot_ui, txt])
            txt.submit(chatbot, [txt, temperature, top_p, max_tokens], [chatbot_ui, txt])
            btn_clear.click(clear_history, [], [chatbot_ui, txt])

        with gr.Column(scale=2):
            with gr.Group(elem_classes="card"):
                gr.Markdown("### 💡 Sugestões rápidas")
                for q in SUGGESTION_QUESTIONS:
                    gr.Button(q, elem_classes="suggestion-btn").click(lambda s=q: s, outputs=[txt])

                gr.Markdown("---")
                gr.Markdown("### 📚 Explorar por tema")

                for theme, qs in SUGGESTIONS_THEMES.items():
                    with gr.Accordion(theme, open=False):
                        for q in qs:
                            gr.Button(q, elem_classes="suggestion-btn").click(lambda s=q: s, outputs=[txt])

    gr.Markdown('<div class="app-footer">EcoLexIA · Sistema RAG para legislação ambiental em Portugal</div>')


if __name__ == "__main__":
    demo.queue()  # recomendado para streaming/requests no Spaces
    demo.launch(
        theme=gr.themes.Soft(),
        css=custom_css,
        ssr_mode=False,   # DESLIGA SSR experimental (evita esse erro do asyncio)
    )