#!/usr/bin/env python

import os
from collections.abc import Iterator
from threading import Thread

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

DESCRIPTION = " "

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

if torch.cuda.is_available():
    model_id = "mistralai/Mistral-7B-Instruct-v0.3"
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_id)


@spaces.GPU
def generate(message: str, chat_history: list[dict]) -> Iterator[str]:
    """Fonction de génération sans sliders : les paramètres
    de génération (max_new_tokens, température, etc.) sont
    fixés en dur.
    """
    # On convertit la question en minuscules pour la comparer plus facilement
    lower_msg = message.strip().lower()

    # Réponse fixe pour "C'est quoi un viol ?"
    if lower_msg == "c’est quoi un viol ?" or lower_msg == "c'est quoi un viol ?":
        fixed_response = """
<h2 style="text-align: center;">Un viol</h2>

<div style="border: 2px solid #fff; padding: 10px; margin: 10px 0;">
  <p>Le viol est un <strong>crime</strong> <br/>
  (donc une infraction très grave punie d'une peine d'emprisonnement)</p>
</div>

<div style="border: 2px solid #fff; padding: 10px; margin: 10px 0;">
<p>Le viol se définit comme :</p>
<ul>
  <li>Tous actes de pénétration sexuelle avec un sexe ou un objet</li>
  <p style="text-align: center;">OU</p>
  <li>Toute pratique sexuelle mettant en contact la bouche d’un partenaire et le sexe d’un autre</li>
</ul>
</div>

<div style="display: flex; align-items: flex-start; justify-content: center; gap: 20px;">

  <!-- Bloc gauche -->
  <div style="border: 2px solid #fff; padding: 10px; margin: 10px 0; width: 40%;">
    <p><strong>SOIT entre</strong> :</p>
    <ul>
      <li><strong>Un majeur</strong> (18 ans ou plus) et un mineur âgé de moins de quinze ans</li>
      <li>Avec une <strong>différence d'âge d'au moins cinq ans</strong> entre les deux personnes (exemple : tu as 13 ans, l’autre a 18 ans)</li>
    </ul>
  </div>

  <!-- Texte "OU" au milieu -->
  <p style="font-weight: bold; margin-top: 20px;">OU</p>

  <!-- Bloc droit -->
  <div style="border: 2px solid #fff; padding: 10px; margin: 10px 0; width: 40%;">
    <p><strong>SOIT avec</strong> :</p>
    <ul>
      <li><strong>Une menace</strong> (ex : chantage)</li>
      <li><strong>Une contrainte</strong> (ex : pression physique)</li>
      <li><strong>Une surprise</strong> (ex : ton partenaire dort)</li>
    </ul>
  </div>

</div>

<div style="border: 2px solid #fff; padding: 10px; margin: 10px 0;">
<p><strong>Ainsi, le viol peut être commis sur</strong> :</p>
<ul>
  <li>La personne qui pénètre</li>
  <li>La personne qui est pénétrée</li>
  <li>La personne qui n’a pas consenti ou ne consent plus</li>
</ul
</div>

<p style="text-align: center;"><em>👉Qu’est-ce que je fais si j’ai vécu ça ?  ?</em> Retourne au menu principal et clique sur le nuage « j’ai besoin de ton aide ».</p>
<p>👉Comment est puni.e le.a violeur ? <p/>
<p>🔎Retrouves la définition sur le site Légifrance, Articles 222-22 et suivants du Code pénal</p>
"""
        yield fixed_response
        return

    # (Exemple) Réponse fixe pour "C'est quoi le consentement ? Comment savoir si ma copine a envie de moi ?"
    if lower_msg == "c’est quoi le consentement ? comment savoir si ma copine a envie de moi ?":
        fixed_response = """Le consentement, pour exister, doit passer par : 

**Le questionnement de soi-même** : il semble nécessaire de s'intéresser à ce que l’on veut précisément et ce que l’on ne veut pas. Savoir clairement ce que l’on veut faire permet d’être en capacité de le dire à l’autre à haute voix.

Le questionnement de l’autre sur ses envies (clique sur le nuage comment demander le consentement de l’autre si tu veux plus de détails).

Le respect de son corps et du corps de l’autre, ainsi que de ses envies : faire quelque chose qui n’est pas expressément désiré par l’autre, tu ne feras point.

⚠️ **Attention** : Le consentement se limite à un moment donné. Tu as le droit de ne pas consentir à une pratique spécifique, de ne plus consentir au rapport sexuel même après avoir donné ton consentement. Tu es précieux. Ainsi, il n’existe pas de devoir sexuel : aucun acte n’est dû à l’autre dans toutes les situations. Le consentement se périme. Il n’existe que de manière temporaire et pour une action particulière.

<h2 style="text-align: center;">Mise en situation</h2>
<p>Billy et Alex s’embrassent. Alex glisse ses mains sous le pull de Billy. Mais Billy n’est pas sûr d’en avoir envie et se sent mal à l’aise. Billy a l’impression d’être touché sans son aval. Alex aurait dû lui demander “Est-ce que je peux ?” ou “Est-ce que tu en as envie ?”</p>
<p>Billy arrête d’embrasser Alex et lui dit “J’ai envie de toi, mais je veux aller à mon rythme.” Cela ne veut pas dire que Billy refusera pour toujours tout contact avec Alex. Dans un monde normal, Billy et Alex se demandent à chaque progression de leur rapport s’ils ont envie.</p>

*<p style="text-align: center;">👉 Comment s’assurer de son consentement ?</p>*
*<p>Retournes sur la page principale et cliques sur le nuage correspondant pour voir toutes les façons de faire. Choisis ce qui te convient et ce qui te rend à l’aise 😌</p>*
"""
        yield fixed_response
        return

    # Sinon, génération classique
    max_new_tokens = 1024
    temperature = 0.6
    top_p = 0.9
    top_k = 50
    repetition_penalty = 1.2

    conversation = [*chat_history, {"role": "user", "content": message}]
    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")

    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)


# Interface Gradio
demo = gr.ChatInterface(
    fn=generate,
    stop_btn=None,
    examples=[
        ["C’est quoi un viol ?"],
        ["C’est quoi le consentement ? Comment savoir si ma copine a envie de moi ?"],
        ["C’est quoi une agression sexuelle ?"],
        ["C’est quoi un attouchement ?"],
        ["C’est quoi un harcèlement sexuel ?"],
        ["Est-ce illégal de visionner du porno ?"],
        ["Mon copain me demande un nude, dois-je le faire ?"],
        ["Mon ancien copain me menace de poster des photos de moi nue sur internet, que faire ?"],
        ["Que puis-je faire si un membre de ma famille me touche d’une manière bizarre, mais que j’ai peur de parler ou de ne pas être cru ?"],
    ],
    type="messages",
    description=DESCRIPTION,
    css_paths="style.css",
)

if __name__ == "__main__":
    demo.queue(max_size=20).launch(share=True)