app.py

"""
NormAssist — Interface Gradio pour l'assistant documentaire technique.

Lancer avec : python app.py
"""

import logging
from pathlib import Path

import gradio as gr

from src.config import DATA_DIR
from src.rag_engine import RAGEngine

# ── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s │ %(name)-25s │ %(levelname)-7s │ %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("normassist")

# ── Initialisation du moteur RAG ─────────────────────────────────────────────

# Auto-générer les PDFs de démo s'ils n'existent pas (utile pour HF Spaces)
sample_dir = DATA_DIR / "sample"
if not list(sample_dir.glob("*.pdf")):
    logger.info("Génération des documents de démonstration...")
    try:
        from scripts.generate_samples import main as generate_samples
        generate_samples()
    except Exception as e:
        logger.warning("Impossible de générer les PDFs de démo : %s", e)

engine = RAGEngine()


# ── Handlers Gradio ──────────────────────────────────────────────────────────


def handle_upload(files) -> str:
    """Gère l'upload d'un ou plusieurs PDFs."""
    if not files:
        return "⚠️ Aucun fichier sélectionné."

    results = []
    for file in files:
        file_path = file.name if hasattr(file, "name") else str(file)
        try:
            msg = engine.ingest_pdf(file_path)
            results.append(msg)
        except Exception as e:
            results.append(f"❌ Erreur : {e}")

    return "\n".join(results)


def handle_question(question: str, history: list) -> tuple:
    """Traite une question et retourne la réponse + sources."""
    if not question.strip():
        return "", history, ""

    result = engine.ask(question)
    answer = result["answer"]

    # Formater les sources pour l'affichage
    sources_text = _format_sources_display(result["sources"])

    # Ajouter à l'historique du chat
    history = history or []
    history.append({"role": "user", "content": question})
    history.append({"role": "assistant", "content": answer})    

    return "", history, sources_text


def handle_load_sample() -> str:
    """Charge les documents de démonstration."""
    sample_dir = DATA_DIR / "sample"
    if not sample_dir.exists() or not list(sample_dir.glob("*.pdf")):
        return (
            "📂 Aucun document de démonstration trouvé.\n\n"
            "Placez des PDFs dans `data/sample/` et relancez.\n\n"
            "Suggestions :\n"
            "• Fiches techniques Schneider Electric (schneider-electric.com)\n"
            "• Datasheets Arduino (arduino.cc)\n"
            "• Documentation ROS (docs.ros.org)\n"
            "• Normes ISO publiques (iso.org/obp)"
        )
    results = engine.ingest_directory(sample_dir)
    return "\n".join(results)


def get_status() -> str:
    """Retourne le statut du système."""
    status = engine.status
    docs = status["liste_documents"]
    
    # On utilise des balises <br> pour les retours à la ligne en HTML
    doc_list = "<br>".join(f"&nbsp;&nbsp;• {d}" for d in docs) if docs else "&nbsp;&nbsp;(aucun)"

    # On force les couleurs avec les variables Gradio et on recrée ta jolie barre bleue à gauche
    return f"""
    <div style="color: var(--body-text-color); background-color: var(--background-fill-secondary); padding: 15px; border-radius: 8px; border-left: 5px solid #3b82f6;">
        <h3 style="margin-top: 0; color: var(--body-text-color);">📊 État du système</h3>
        <b>Documents indexés :</b> {status['documents_indexés']}<br>
        {doc_list}<br><br>
        <b>Chunks en base :</b> {status['chunks_total']}<br>
        <b>Vecteurs FAISS :</b> {status['vecteurs_total']}<br>
        <b>LLM :</b> {status['modèle_llm']}
    </div>
    """


def _format_sources_display(sources: list[dict]) -> str:
    """Formate les sources pour l'affichage dans le panneau latéral."""
    if not sources:
        return "Aucune source trouvée."

    parts = []
    for i, s in enumerate(sources, 1):
        parts.append(
            f"**Source {i}** — {s['document']}, page {s['page']} "
            f"(score: {s['score']})\n\n"
            f"> {s['extrait']}\n"
        )
    return "\n---\n".join(parts)


# ── Thème et CSS ─────────────────────────────────────────────────────────────

CUSTOM_CSS = """
.main-title {
    text-align: center;
    margin-bottom: 0.5rem;
}
.main-title h1 {
    font-size: 2rem;
    font-weight: 700;
    color: #1a1a2e;
}
.subtitle {
    text-align: center;
    color: #666;
    font-size: 0.95rem;
    margin-bottom: 1.5rem;
}
.status-box {
    background: #f8f9fa;
    border-left: 4px solid #4a90d9;
    padding: 1rem;
    border-radius: 0 8px 8px 0;
    font-size: 0.9rem;
}
"""

# ── Interface Gradio ─────────────────────────────────────────────────────────


def build_app() -> gr.Blocks:
    """Construit l'interface Gradio."""

    with gr.Blocks(
        title="NormAssist — Documentation Technique IA",
        css=CUSTOM_CSS,
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="slate",
            font=gr.themes.GoogleFont("Source Sans Pro"),
        ),
    ) as app:
        # ── En-tête ──────────────────────────────────────────────────────
        gr.HTML(
            '<div class="main-title"><h1>🏭 NormAssist</h1></div>'
            '<div class="subtitle">'
            "Assistant IA pour documentation technique industrielle — "
            "Normes ISO · Fiches techniques · Manuels de maintenance"
            "</div>"
        )

        with gr.Row():
            # ── Colonne gauche : gestion documentaire ────────────────────
            with gr.Column(scale=1):
                gr.Markdown("### 📁 Base documentaire")

                upload = gr.File(
                    label="Uploader des PDFs",
                    file_count="multiple",
                    file_types=[".pdf"],
                    type="filepath",
                )
                upload_btn = gr.Button("📥 Indexer les documents", variant="primary")
                upload_result = gr.Markdown(label="Résultat")

                gr.Markdown("---")

                sample_btn = gr.Button("📂 Charger les documents de démo")
                sample_result = gr.Markdown()

                gr.Markdown("---")

                status_btn = gr.Button("📊 Voir le statut")
                status_display = gr.Markdown(elem_classes=["status-box"])

            # ── Colonne centrale : chat ──────────────────────────────────
            with gr.Column(scale=2):
                gr.Markdown("### 💬 Posez votre question")

                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=450,
                    show_label=False,
                    avatar_images=(None, "🏭"),
                )

                with gr.Row():
                    question_input = gr.Textbox(
                        placeholder=(
                            "Ex: Quelles sont les exigences de traçabilité "
                            "dans cette norme ?"
                        ),
                        label="Question",
                        scale=5,
                        show_label=False,
                    )
                    ask_btn = gr.Button("🔍 Chercher", variant="primary", scale=1)

                gr.Markdown("---")

                gr.Markdown("### 📎 Sources utilisées")
                sources_display = gr.Markdown(
                    value="Les sources apparaîtront ici après une question."
                )

        # ── Exemples ─────────────────────────────────────────────────────
        gr.Markdown("### 💡 Exemples de questions")
        gr.Examples(
            examples=[
                "Quelles sont les exigences de traçabilité dans ce document ?",
                "Quelle est la température maximale d'utilisation ?",
                "Décris la procédure de maintenance préventive.",
                "Quels sont les critères d'acceptation définis ?",
                "Liste les équipements de protection individuelle requis.",
            ],
            inputs=question_input,
        )

        # ── Events ───────────────────────────────────────────────────────
        upload_btn.click(
            fn=handle_upload, inputs=[upload], outputs=[upload_result]
        )

        sample_btn.click(fn=handle_load_sample, outputs=[sample_result])

        status_btn.click(fn=get_status, outputs=[status_display])

        # Question via bouton ou Entrée
        ask_btn.click(
            fn=handle_question,
            inputs=[question_input, chatbot],
            outputs=[question_input, chatbot, sources_display],
        )
        question_input.submit(
            fn=handle_question,
            inputs=[question_input, chatbot],
            outputs=[question_input, chatbot, sources_display],
        )

    return app


# ── Point d'entrée ───────────────────────────────────────────────────────────

if __name__ == "__main__":
    app = build_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
    )