from __future__ import annotations

try:
    from .taxonomy import JIGSAW_LABELS, SYNTHETIC_ONLY, TAXONOMY
except ImportError:
    from taxonomy import JIGSAW_LABELS, SYNTHETIC_ONLY, TAXONOMY
CIVIL = "google/civil_comments"
JIGSAW = "google/jigsaw_toxicity_pred"
JIGSAW_MIRROR = "tasksource/jigsaw_toxicity"
JIGSAW_EVAL_MOD = 20
MHS = "ucberkeley-dlab/measuring-hate-speech"
RTP = "allenai/real-toxicity-prompts"


def jigsaw_is_eval(row):
    import zlib

    return zlib.crc32(str(row.get("id", "")).encode()) % JIGSAW_EVAL_MOD == 0


TAXONOMY_SOURCE_MAP = {
    "harassment": [
        (CIVIL, "insult", 0.5),
        (JIGSAW, "insult", 1),
        (MHS, "insult", 3),
        (MHS, "humiliate", 3),
    ],
    "harassment_threatening": [
        (CIVIL, ("threat", "insult"), 0.5),
        (JIGSAW, ("threat", "insult"), 1),
        (MHS, ("violence", "insult"), 3),
    ],
    "hate": [
        (CIVIL, "identity_attack", 0.5),
        (JIGSAW, "identity_hate", 1),
        (MHS, "hate_speech_score", 0.5),
    ],
    "hate_threatening": [
        (CIVIL, ("identity_attack", "threat"), 0.5),
        (JIGSAW, ("identity_hate", "threat"), 1),
        (MHS, ("hate_speech_score", "violence"), (0.5, 3)),
    ],
    "self_harm": [],
    "self_harm_intent": [],
    "self_harm_instructions": [],
    "sexual": [
        (CIVIL, "sexual_explicit", 0.5),
        (RTP, "sexually_explicit", 0.5),
        (RTP, "flirtation", 0.5),
    ],
    "sexual_minors": [],
    "violence": [(MHS, "violence", 2), (CIVIL, "threat", 0.5), (JIGSAW, "threat", 1)],
    "violence_graphic": [],
    "identity_attack": [
        (CIVIL, "identity_attack", 0.5),
        (MHS, "hate_speech_score", 0.5),
        (RTP, "identity_attack", 0.5),
    ],
    "insult": [(CIVIL, "insult", 0.5), (JIGSAW, "insult", 1), (MHS, "insult", 3)],
    "profanity_obscene": [
        (CIVIL, "obscene", 0.5),
        (JIGSAW, "obscene", 1),
        (RTP, "profanity", 0.5),
    ],
    "threat": [(CIVIL, "threat", 0.5), (JIGSAW, "threat", 1), (RTP, "threat", 0.5)],
    "sexual_explicit": [
        (CIVIL, "sexual_explicit", 0.5),
        (RTP, "sexually_explicit", 0.5),
    ],
    "severe_toxicity": [(CIVIL, "severe_toxicity", 0.5), (JIGSAW, "severe_toxic", 1)],
    "dehumanize": [(MHS, "dehumanize", 3)],
    "incitement_violence": [(MHS, "violence", 3)],
}
JIGSAW_SOURCE_MAP = {
    "toxic": [(JIGSAW, "toxic", 1), (CIVIL, "toxicity", 0.5)],
    "severe_toxic": [(JIGSAW, "severe_toxic", 1), (CIVIL, "severe_toxicity", 0.5)],
    "obscene": [(JIGSAW, "obscene", 1), (CIVIL, "obscene", 0.5)],
    "threat": [(JIGSAW, "threat", 1), (CIVIL, "threat", 0.5)],
    "insult": [(JIGSAW, "insult", 1), (CIVIL, "insult", 0.5)],
    "identity_hate": [(JIGSAW, "identity_hate", 1), (CIVIL, "identity_attack", 0.5)],
}
MASKED_LABELS = [lbl for lbl in TAXONOMY if not TAXONOMY_SOURCE_MAP.get(lbl)]
DATA_BACKED_LABELS = [lbl for lbl in TAXONOMY if TAXONOMY_SOURCE_MAP.get(lbl)]
assert set(SYNTHETIC_ONLY).issubset(set(MASKED_LABELS)), (
    "SYNTHETIC_ONLY labels must have no commercial source rule"
)


def licenses():
    return [
        (CIVIL, "CC0-1.0", False),
        (JIGSAW, "CC0-1.0", False),
        (MHS, "CC-BY-4.0", True),
        (RTP, "Apache-2.0", False),
    ]


def _eval_rule(rule, row):
    _dataset, column, threshold = rule
    cols = column if isinstance(column, tuple) else (column,)
    thrs = threshold if isinstance(threshold, tuple) else (threshold,) * len(cols)
    fired = True
    for col, thr in zip(cols, thrs):
        val = row.get(col)
        if val is None:
            return None
        try:
            num = float(val)
        except (TypeError, ValueError):
            return None
        if num < thr:
            fired = False
    return fired


def _row_targets(source_map, index_of, n_labels, dataset_id, row):
    labels = [0.0] * n_labels
    mask = [0.0] * n_labels
    for label, rules in source_map.items():
        result = None
        for rule in rules:
            if rule[0] != dataset_id:
                continue
            fired = _eval_rule(rule, row)
            if fired is None:
                continue
            result = bool(result) or fired
        if result is None:
            continue
        idx = index_of[label]
        mask[idx] = 1.0
        labels[idx] = 1.0 if result else 0.0
    return (labels, mask)


def _iter_civil(streaming, max_rows):
    from datasets import load_dataset

    ds = load_dataset(CIVIL, split="train", streaming=streaming)
    for i, ex in enumerate(ds):
        if max_rows is not None and i >= max_rows:
            break
        text = ex.get("text") or ""
        if text:
            yield (text, ex)


def _iter_mhs(streaming, max_rows):
    from datasets import load_dataset

    ds = load_dataset(MHS, split="train", streaming=streaming)
    for i, ex in enumerate(ds):
        if max_rows is not None and i >= max_rows:
            break
        text = ex.get("text") or ""
        if text:
            yield (text, ex)


def _iter_rtp(streaming, max_rows):
    from datasets import load_dataset

    ds = load_dataset(RTP, split="train", streaming=streaming)
    count = 0
    for ex in ds:
        for key in ("prompt", "continuation"):
            part = ex.get(key) or {}
            text = part.get("text") or ""
            if not text:
                continue
            if max_rows is not None and count >= max_rows:
                return
            count += 1
            yield (text, part)


def _iter_jigsaw(streaming, max_rows, want_eval=False):
    import os
    from datasets import load_dataset

    jig_dir = os.environ.get("JIGSAW_DIR")
    try:
        if jig_dir:
            ds = load_dataset(
                JIGSAW, split="train", streaming=streaming, data_dir=jig_dir
            )
        else:
            ds = load_dataset(JIGSAW_MIRROR, split="train", streaming=streaming)
        kept = 0
        for ex in ds:
            if jigsaw_is_eval(ex) != want_eval:
                continue
            if max_rows is not None and kept >= max_rows:
                break
            text = ex.get("comment_text") or ""
            if text:
                kept += 1
                yield (text, ex)
    except Exception as err:
        print(
            f"[toxicity] real Jigsaw unavailable ({err}); falling back to the CC0 Civil Comments backbone, which shares the Jigsaw-head columns so nothing trains on 0s."
        )
        return


def load_taxonomy_examples(max_rows=None, streaming=True):
    index_of = {lbl: i for i, lbl in enumerate(TAXONOMY)}
    n_labels = len(TAXONOMY)
    sources = [
        (CIVIL, _iter_civil),
        (MHS, _iter_mhs),
        (RTP, _iter_rtp),
        (JIGSAW, _iter_jigsaw),
    ]
    for dataset_id, iterator in sources:
        for text, row in iterator(streaming, max_rows):
            labels, mask = _row_targets(
                TAXONOMY_SOURCE_MAP, index_of, n_labels, dataset_id, row
            )
            if sum(mask) == 0:
                continue
            yield {"text": text, "labels": labels, "mask": mask}


def load_jigsaw_examples(max_rows=None, streaming=True):
    index_of = {lbl: i for i, lbl in enumerate(JIGSAW_LABELS)}
    n_labels = len(JIGSAW_LABELS)
    sources = [(JIGSAW, _iter_jigsaw), (CIVIL, _iter_civil)]
    for dataset_id, iterator in sources:
        for text, row in iterator(streaming, max_rows):
            labels, mask = _row_targets(
                JIGSAW_SOURCE_MAP, index_of, n_labels, dataset_id, row
            )
            if sum(mask) == 0:
                continue
            yield {"text": text, "labels": labels, "mask": mask}