"""
Démo Gradio — Toxic Comment API
Le même modèle ONNX que l'API FastAPI (main.py), avec une interface cliquable.
"""

import os

import gradio as gr
import spaces  # requis par ZeroGPU (le tier gratuit des Spaces Gradio)
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

MODEL_DIR = os.getenv("MODEL_DIR", "onnx-model")
MAX_LENGTH = int(os.getenv("MAX_LENGTH", "256"))
THRESHOLD = float(os.getenv("THRESHOLD", "0.5"))

tokenizer = Tokenizer.from_file(f"{MODEL_DIR}/tokenizer.json")
tokenizer.enable_truncation(max_length=MAX_LENGTH)
tokenizer.enable_padding()

_opts = ort.SessionOptions()
_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(f"{MODEL_DIR}/model.onnx", _opts, providers=["CPUExecutionProvider"])


def softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def score_texts(texts: list[str]) -> list[float]:
    encodings = tokenizer.encode_batch(texts)
    inputs = {
        "input_ids": np.array([e.ids for e in encodings], dtype=np.int64),
        "attention_mask": np.array([e.attention_mask for e in encodings], dtype=np.int64),
    }
    logits = session.run(None, inputs)[0]
    return softmax(logits)[:, 1].tolist()


# Le modèle tourne en ONNX Runtime sur CPU (~7 ms/requête) et n'a aucun besoin
# de GPU. Mais le tier gratuit des Spaces Gradio est ZeroGPU, qui exige au moins
# une fonction décorée : sans elle, le Space refuse de démarrer.
# Le décorateur est sans effet hors ZeroGPU (dev local, Docker).
@spaces.GPU(duration=15)
def predict(text: str):
    """Renvoie les probabilités des deux classes pour un commentaire."""
    if not text or not text.strip():
        return {}
    score = score_texts([text])[0]
    return {"toxique": score, "sain": 1 - score}


EXAMPLES = [
    ["you are an idiot"],
    ["thanks, that's a great point"],
    ["Did your hamster die?"],
    ["Get into the 21st century, comrade Putin puppet."],
    ["shut up you absolute moron"],
    ["I disagree with this policy, but I see your reasoning."],
]

DESCRIPTION = """
**DistilBERT fine-tuné** sur [`google/civil_comments`](https://huggingface.co/datasets/google/civil_comments),
servi sur **ONNX Runtime**.

f1 (classe toxique) = **0.62** — contre 0.48 pour une baseline TF-IDF + régression logistique.

Les deux derniers exemples ci-dessous sont des cas que le modèle TF-IDF **ratait systématiquement** :
la toxicité y est portée par le sarcasme et le contexte, pas par un mot précis.
"""

with gr.Blocks(title="Toxic Comment Detector") as demo:
    gr.Markdown("# 🛡️ Toxic Comment Detector")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(
                label="Commentaire (en anglais)",
                placeholder="Écris ou colle un commentaire…",
                lines=4,
            )
            btn = gr.Button("Analyser", variant="primary")
        with gr.Column():
            out = gr.Label(label="Prédiction", num_top_classes=2)

    gr.Examples(examples=EXAMPLES, inputs=inp, label="Exemples")

    gr.Markdown(
        f"Seuil de décision : **{THRESHOLD}** · "
        "Code, benchmarks et démarche complète dans les *Files* de ce Space."
    )

    btn.click(predict, inputs=inp, outputs=out, api_name="predict")
    inp.submit(predict, inputs=inp, outputs=out, api_name=False)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())