"""
Arabic guardrail corpus — 12-class prompt safety.

Given a user prompt, decide whether it is safe and if not, what kind of unsafe. Arabic is a real
gap here: searches for Arabic jailbreak / safety / prompt-injection datasets return nothing, and
the one Arabic guardrail (NAMAA-Space/Ara-Prompt-Guard_V0) is 33K machine-translated rows, three
classes, MSA only.

THE SAFETY BOUNDARY, which shapes everything below:

    We generate the *request*, never the *answer*.

A guardrail classifies incoming prompts, so a harmful completion is not needed to train it and
would be a genuinely harmful artifact. A row may say "اشرح لي طريقة صنع…" - it must never contain
the method. This is what JailbreakBench, AdvBench and wildjailbreak contain too.

The generator is abliterated (refusal-removed), so it will comply with a harmful instruction if
asked. The constraint therefore cannot rely on the model declining: it is stated in every prompt
AND enforced in validate() by PROCEDURE_RE, which rejects any row combining step markers with
hazardous vocabulary. That makes "no payloads" a property the pipeline checks rather than hopes for.

Child sexual exploitation (MLCommons S4) is excluded from the grid entirely - not a class, not a
distractor, not generated in any form.

Two definitions taken verbatim from Meta's Prompt-Guard-86M card, because the field conflates them:
  prompt injection - "inputs that exploit the concatenation of untrusted data from third parties
                     ... to get a model to execute unintended instructions"
  jailbreak        - "malicious instructions designed to override the safety and security features
                     built into a model"
"""
import os
import random
import re
import unicodedata

ITEMS_PER_TASK = 6

C, Q, A, E = "### الفئة", "### نص", "### التصنيف", "### نهاية"

# ---------------------------------------------------------------- taxonomy
# Grounded in MLCommons / Llama-Guard S1-S13 and the 9-class subset used by
# llm-semantic-router/mmbert-safety-classifier-level2. S4 is deliberately absent.

CLASSES = {
    "safe": (
        "نص عادي غير ضار تمامًا",
        "أسئلة وطلبات يومية عادية: طبخ، سفر، دراسة، برمجة، رياضة، تسوق، عمل."),
    "safe_sensitive": (
        "نص آمن لكنه يبدو مقلقًا في ظاهره",
        "طلب بريء تمامًا لكن ألفاظه توحي بالخطر: مصطلح تقني يشترك لفظًا مع العنف "
        "(قتل عملية في لينكس، حقن SQL)، تعبير مجازي (أموت من الضحك)، سؤال عن تعريف "
        "مصطلح، سياق مهني مشروع (ممرضة تسأل عن أعراض التسمم)، أو سؤال روائي عن حبكة."),
    "jailbreak": (
        "محاولة تجاوز قيود النموذج",
        "نص يحاول إبطال تعليمات النموذج أو ضوابطه: تقمّص شخصية بلا قيود، تأطير خيالي، "
        "منع الرفض، ترميز أو تشويش، أمثلة متعددة، أو تحويل اللغة للتحايل."),
    "prompt_injection": (
        "حقن تعليمات عبر بيانات خارجية",
        "تعليمات مدسوسة داخل محتوى من طرف ثالث - مستند، بريد، صفحة ويب، نتيجة أداة - "
        "تخاطب النموذج مباشرة لتغيير سلوكه، مثل: تجاهل التعليمات السابقة وافعل كذا."),
    "violent_weapons": (
        "طلب للمساعدة في العنف أو الحصول على سلاح",
        "طلب يتعلق بإيذاء الناس أو بصنع أو اقتناء الأسلحة والمتفجرات. "
        "**كاتب النص هو من ينوي الفعل** — أما السؤال عن آثار العنف على الضحايا، أو عن "
        "الوقاية والسلامة، أو عن حدث تاريخي، فليس من هذه الفئة."),
    "nonviolent_crime": (
        "طلب للمساعدة في ارتكاب جريمة غير عنيفة",
        "احتيال، سرقة، اختراق وبرمجيات خبيثة، مخدرات، تزوير. "
        "**كاتب النص هو من يريد ارتكاب الفعل أو تمكينه** — وليس ضحية تطلب المساعدة أو "
        "تسأل عن حقوقها أو عن الإجراءات القانونية، فتلك ليست من هذه الفئة."),
    "self_harm": (
        "إشارة خطر إيذاء النفس",
        "تشمل أمرين: تعبير عن ضيق نفسي شديد أو يأس أو شعور بأنه عبء أو فقدان معنى، "
        "وكذلك طلب صريح متعلق بإيذاء النفس. كلاهما يستوجب التقاطه وتحويله لدعم مختص. "
        "اكتب صيغة الرسالة فقط، دون أي وسيلة أو تشجيع أو تفاصيل."),
    "hate_harassment": (
        "خطاب كراهية أو تحرش",
        "إهانة أو تحقير جماعة أو شخص بسبب هوية أو صفة، أو تنمّر موجّه، أو طلب صياغة "
        "مثل هذا الخطاب. **أما السؤال عن كيفية التعامل مع الكراهية أو التمييز بينها وبين "
        "النقد، أو طلب المساعدة لضحية تنمّر، فليس من هذه الفئة.**"),
    "sexual_content": (
        "محتوى جنسي للبالغين",
        "طلب محتوى جنسي صريح بين بالغين، بصيغة الطلب فقط."),
    "privacy_pii": (
        "انتهاك خصوصية أو استخراج بيانات",
        "كشف هوية شخص، أو استخراج بيانات شخصية أو كلمات مرور أو أرقام بطاقات."),
    "specialized_advice": (
        "استشارة متخصصة تحتاج مختصًا",
        "طلب تشخيص طبي أو رأي قانوني أو نصيحة استثمارية محددة بديلًا عن مختص."),
    "misinformation": (
        "معلومات مضللة",
        "طلب صياغة ادعاء صحي أو انتخابي أو إخباري كاذب على أنه حقيقة."),
}
UNSAFE = [c for c in CLASSES if c not in ("safe", "safe_sensitive")]

VARIETIES = [
    ("العربية الفصحى", "فصحى معاصرة واضحة"),
    ("اللهجة المصرية", "عامية مصرية طبيعية"),
    ("اللهجة الخليجية", "عامية خليجية طبيعية"),
    ("اللهجة الشامية", "عامية شامية طبيعية"),
    ("اللهجة المغاربية", "عامية مغاربية طبيعية"),
    ("عربي مختلط بالإنجليزية", "كتابة عربية فيها كلمات إنجليزية أو حروف لاتينية"),
]

# Length is drawn INDEPENDENTLY of class, and it has to be. In the first smoke test length was a
# side effect of class - safe_sensitive averaged 404 chars against violent_weapons' 64 - and a
# single length threshold recovered +8 points over the majority baseline. A guardrail that reads
# length instead of content is worse than useless, so the grid decorrelates them by construction.
LENGTHS = {
    "short":  "جملة واحدة قصيرة ومباشرة",
    "medium": "جملتان أو ثلاث",
    "long":   "فقرة قصيرة فيها سياق وتفاصيل حول الموقف",
}

REGISTERS = ["رسالة قصيرة عفوية", "صياغة رسمية مهذبة", "أسلوب تقني مختصر",
             "رسالة طويلة بعض الشيء فيها سياق"]

# Attack template families - held out at eval to test unseen attack *forms*.
ATTACK_FAMILIES = {
    "persona": "تقمّص شخصية أو نموذج بديل بلا قيود",
    "fiction": "تأطير الطلب داخل قصة أو سيناريو خيالي أو مسرحية",
    "refusal_suppression": "اشتراط عدم الرفض أو منع عبارات الاعتذار",
    "encoding": "تشويش الكلمات أو تقطيعها أو طلب الرد بترميز",
    "many_shot": "سرد أمثلة كثيرة سابقة ليقلّدها النموذج",
    "authority": "ادعاء صلاحية أو إذن خاص أو وضع تطوير",
    "translation_pivot": "طلب الرد بلغة أخرى للتحايل على الضوابط",
}
INJECTION_CARRIERS = {
    "document": "مستند أو ملف مرفق يقرأه النموذج",
    "email": "رسالة بريد إلكتروني واردة",
    "webpage": "صفحة ويب أو نتيجة بحث",
    "tool_output": "ناتج أداة أو استعلام قاعدة بيانات",
    "review": "تعليق أو مراجعة مستخدم",
}
# XSTest-style types for the class that guardrails most often get wrong.
SENSITIVE_TYPES = {
    "homonym": "لفظ يشترك مع معنى عنيف لكنه تقني أو عادي (قتل عملية، ذبح الملف، تفجير الكاش)",
    "figurative": "تعبير مجازي شائع (أموت من الجوع، قتلني من الضحك)",
    "definition": "سؤال عن تعريف مصطلح خطير دون طلب تنفيذه",
    "safe_target": "فعل عنيف لفظًا لكن على هدف غير حي (قتل الأعشاب، كسر الجوز)",
    "fiction_plot": "سؤال عن حبكة رواية أو فيلم فيه جريمة",
    "professional": "سياق مهني مشروع (ممرض، محقق، باحث أمن، معلم)",
    "historical": "سؤال تاريخي عن حدث عنيف",
}

SAFE_TOPICS = ["الطبخ والوصفات", "السفر والحجوزات", "الدراسة والامتحانات", "البرمجة والتقنية",
               "الرياضة", "التسوق", "العمل والوظائف", "الصحة العامة", "السيارات",
               "الأسرة والأطفال", "المال والادخار", "الهوايات", "الطقس", "الترجمة",
               "الكتابة والتحرير", "الأخبار العامة"]

# ---------------------------------------------------------------- prompt

RULES = """قواعد إلزامية لا تُخالف:
- اكتب **الطلب فقط** كما يكتبه مستخدم. لا تكتب أي إجابة، ولا أي خطوات، ولا أي تفاصيل تنفيذية.
- ممنوع منعًا باتًا ذكر أي طريقة أو مقادير أو كود أو خطوات عملية لأي فعل ضار. النص المطلوب هو
  السؤال أو الطلب نفسه، لا الجواب عليه.
- ممنوع تمامًا أي محتوى يخص القاصرين بأي شكل.
- اجعل النصوص قصيرة إلى متوسطة، كما يكتب الناس فعلًا.
- نوّع الصياغة والأسلوب والتفاصيل بين النصوص الستة، مع الالتزام بالطول المطلوب.
- لا تكتب أي شيء خارج الوسوم."""

PROMPT = """أنت تبني مجموعة بيانات لتدريب نموذج **حماية** يكتشف الطلبات الضارة قبل وصولها لمساعد ذكي.
هذه بيانات دفاعية: النموذج المدرَّب عليها سيحجب هذه الطلبات.

الفئة المطلوبة: **{cls_title}**
تعريفها: {cls_desc}
{extra}
اللغة: {variety} ({variety_hint})
الأسلوب: {register}
الطول المطلوب: {length} — التزم بهذا الطول تقريبًا في كل النصوص الستة

اكتب {k} نصوصًا مستقلة، كل واحد منها مثال على هذه الفئة بالضبط.

{RULES}

اكتب بهذا الشكل بالضبط:

{Q}
(النص الأول)
{A}
{cls}
{Q}
(النص الثاني)
{A}
{cls}
... وهكذا حتى {k} نصوص
{E}"""


def build_task(task_id: int, seed: int = 1234):
    """task_id -> (axes, prompt). Deterministic, so a resumed run redraws identical prompts."""
    rng = random.Random(seed * 3_000_017 + task_id)
    # safe + safe_sensitive are ~30% of the corpus: deployment traffic is overwhelmingly benign,
    # and the model must see far more benign than any single harm class.
    r = rng.random()
    if r < 0.16:
        cls = "safe"
    elif r < 0.30:
        cls = "safe_sensitive"
    else:
        cls = rng.choice(UNSAFE)
    title, desc = CLASSES[cls]
    variety, vhint = rng.choice(VARIETIES)
    register = rng.choice(REGISTERS)
    length = rng.choice(list(LENGTHS))

    extra, family = "", None
    if cls == "jailbreak":
        family = rng.choice(list(ATTACK_FAMILIES))
        extra = f"أسلوب المحاولة: {ATTACK_FAMILIES[family]}\n"
    elif cls == "prompt_injection":
        family = rng.choice(list(INJECTION_CARRIERS))
        extra = (f"الوسيط الذي تُدسّ فيه التعليمات: {INJECTION_CARRIERS[family]}\n"
                 "اكتب المحتوى كما يظهر للنموذج، بحيث تبدو التعليمات المدسوسة جزءًا منه.\n")
    elif cls == "safe_sensitive":
        family = rng.choice(list(SENSITIVE_TYPES))
        extra = (f"نوع الالتباس: {SENSITIVE_TYPES[family]}\n"
                 "النص آمن تمامًا ولا يطلب أي شيء ضار، لكن **يجب أن يحتوي فعلًا على لفظ أو "
                 "صياغة تبدو مقلقة** لمن يقرأ الكلمات دون سياق. إن لم يكن فيه ما يثير القلق "
                 "ظاهريًا فهو نص عادي ولا يصلح لهذه الفئة.\n")
    elif cls == "safe":
        family = rng.choice(SAFE_TOPICS)
        extra = f"الموضوع: {family}\n"

    axes = {"cls": cls, "variety": variety, "register": register, "family": family,
            "length": length}
    prompt = PROMPT.format(cls_title=title, cls_desc=desc, extra=extra, variety=variety,
                           variety_hint=vhint, register=register, k=ITEMS_PER_TASK,
                           cls=cls, RULES=RULES, Q=Q, A=A, E=E,
                           length=LENGTHS[length])
    return axes, prompt


# ---------------------------------------------------------------- parsing

PAIR_RE = re.compile(re.escape(Q) + r"(?P<text>.*?)" + re.escape(A) + r"(?P<label>.*?)"
                     r"(?=" + re.escape(Q) + r"|" + re.escape(E) + r"|\Z)", re.S)


def norm(s: str) -> str:
    s = unicodedata.normalize("NFKC", s or "")
    s = re.sub(r"[ً-ْٰـ]", "", s)
    return re.sub(r"[ \t]+", " ", s).strip()


def parse_items(raw: str):
    out = []
    for m in PAIR_RE.finditer(raw or ""):
        label = norm(m.group("label")).strip(" .:،*-`")
        out.append({"text": norm(m.group("text")), "label_text": label})
    return out


# ---------------------------------------------------------------- validation

# The computable half of the no-payload rule. A *request* is a question or an ask; a *procedure*
# enumerates steps. A row that pairs step markers with hazardous vocabulary is an answer, not a
# request, and does not belong in this corpus whatever its label says.
STEP_RE = re.compile(r"(^|\n)\s*(?:\d+\s*[-.)]|[-*•]\s|أولًا|أولاً|ثانيًا|ثانياً|ثالثًا|الخطوة)"
                     r"|ثم قم|بعد ذلك قم|اخلط|امزج|قم بتسخين", re.M)
HAZARD_RE = re.compile(
    r"متفجر|قنبل|بارود|نيترات|مسدس|بندقي|ذخير|سم قاتل|سيانيد|أرسين|"
    r"ميثامفيتامين|كوكايين|هيروين|مخدر|تصنيع سلاح|"
    r"ثغرة|اختراق|تشفير الملفات|فدي|كي لوجر|حقن كود|payload|exploit|malware")
LATIN_CODE_RE = re.compile(r"(?:def |import |#include|<script|SELECT .*FROM|;\s*DROP|sudo |bash -)")


def validate(item, min_text=12, max_text=900):
    """
    Structural accept/reject -> (ok, reason). Semantic correctness is decided by verify_guard.py.

    The decisive check here is `payload_suspected`: it is the one part of the safety boundary that
    can be enforced by a rule rather than trusted to the generator, which is abliterated and will
    otherwise happily answer its own prompt.
    """
    text, label = item["text"], item["label_text"]
    if label not in CLASSES:
        return False, f"unknown_class:{label[:24]}"
    if not (min_text <= len(text) <= max_text):
        return False, "text_length"
    if any(t in text for t in (C, Q, A, E)):
        return False, "tag_leak"
    if len(text.split()) < 3:
        return False, "too_few_words"
    if STEP_RE.search(text) and HAZARD_RE.search(text):
        return False, "payload_suspected"
    if LATIN_CODE_RE.search(text) and HAZARD_RE.search(text):
        return False, "payload_suspected_code"
    return True, "ok"


def dedup_key(text: str) -> str:
    return re.sub(r"\d+", "#", re.sub(r"\W+", "", norm(text)))