"""Clean: drop low-information documents, then remove anything that overlaps the benchmarks you will test on.

The prose rules follow the Gopher and C4 heuristics with FineWeb's thresholds. Math and code sources get their own,
looser rules, because LaTeX and source code fail prose checks by design.
"""
from __future__ import annotations

import re
from collections import Counter
from dataclasses import dataclass, field

from .text import ngrams, norm, words

BOILER = re.compile(r"(cookie|javascript|privacy policy|terms of (use|service)|all rights reserved|lorem ipsum|"
                    r"click here|sign up for our newsletter|accept all|log ?in to (reply|comment)|skip to (main )?content)", re.I)
AUTOGEN = re.compile(r"(auto-?generated|do not edit|generated by (the )?protoc|minified)", re.I)

DEFAULTS = {
    "min_words": 50, "max_words": 100_000, "min_mean_word": 3.0, "max_mean_word": 10.0,
    "max_symbol_ratio": 0.1, "min_alpha_words": 0.8, "max_ellipsis_lines": 0.3,
    "max_dup_lines": 0.3, "max_top2gram": 0.2, "max_dup5gram": 0.15, "max_dup10gram": 0.1, "max_boiler_lines": 0.4,
    "strip_boilerplate_lines": True,
    # code
    "code_max_line": 1000, "code_min_alpha": 0.25,
}


def dup_ngram_chars(ws: list[str], n: int) -> int:
    """Characters inside repeated n-grams, counting each repeat once and never the first copy (Gopher, as
    implemented in datatrove). Overlapping windows are skipped so a repeated paragraph is not counted n times."""
    seen, dup, i = set(), 0, 0
    while i <= len(ws) - n:
        g = tuple(ws[i:i + n])
        if g in seen:
            dup += sum(len(w) for w in g)
            i += n
        else:
            seen.add(g)
            i += 1
    return dup


@dataclass
class Cleaner:
    rules: dict = field(default_factory=lambda: dict(DEFAULTS))

    def __post_init__(self):
        self.rules = {**DEFAULTS, **(self.rules or {})}

    def apply(self, doc: dict) -> tuple[dict | None, str]:
        """Return (document, "") to keep it, possibly with boilerplate lines removed, or (None, reason) to drop it."""
        text = doc["text"]
        kind = doc.get("kind", "web")
        if not text.strip():
            return None, "empty"
        if kind == "code":
            why = self._code(text)
            return (None, why) if why else (doc, "")
        if self.rules["strip_boilerplate_lines"] and kind in ("web", "educational", "multilingual"):
            lines = text.split("\n")
            kept = [l for l in lines if not (BOILER.search(l) and len(l.split()) < 25)]
            if lines and len(lines) - len(kept) > self.rules["max_boiler_lines"] * len(lines):
                return None, "mostly boilerplate"
            if len(kept) != len(lines):
                text = "\n".join(kept)
                doc = {**doc, "text": text}
        if kind == "textbook":  # PDFs repeat running headers and footers on every page; drop those lines, keep the page
            lines = text.split("\n")
            short = Counter(l.strip() for l in lines if l.strip() and len(l.split()) <= 12)
            running = {l for l, c in short.items() if c >= 3}
            if running:
                text = "\n".join(l for l in lines if l.strip() not in running)
                doc = {**doc, "text": text}
        why = self._prose(text, loose=kind in ("math", "textbook", "synthetic"), rep=1.6 if kind == "math" else 1.0)
        return (None, why) if why else (doc, "")

    def _prose(self, text: str, loose: bool, rep: float = 1.0) -> str:
        """rep scales the repetition caps: worked math repeats symbols and steps by nature."""
        r = self.rules
        toks = text.split()
        n = len(toks)
        if n < r["min_words"]:
            return "too short"
        if n > r["max_words"]:
            return "too long"
        if not loose:
            mean = sum(len(t) for t in toks) / n
            if not r["min_mean_word"] <= mean <= r["max_mean_word"]:
                return "odd word lengths"
            sym = (text.count("#") + text.count("...") + text.count("…")) / n
            if sym > r["max_symbol_ratio"]:
                return "too many symbols"
            alpha = sum(1 for t in toks if any(c.isalpha() for c in t)) / n
            if alpha < r["min_alpha_words"]:
                return "too little real text"
        lines = [l.strip() for l in text.split("\n") if l.strip()]
        if lines:
            ell = sum(1 for l in lines if l.endswith(("...", "…"))) / len(lines)
            if ell > r["max_ellipsis_lines"]:
                return "truncated lines"
            if len(lines) >= 5:
                dup = 1 - len(set(lines)) / len(lines)
                if dup > r["max_dup_lines"] * rep:
                    return "repeated lines"
        ws = words(text)
        if len(ws) >= 20:
            bi = Counter(ngrams(ws, 2))
            top = bi.most_common(1)[0]
            if top[1] > 2 and top[1] * len(top[0]) / max(1, sum(len(w) + 1 for w in ws)) > r["max_top2gram"] * rep:
                return "one phrase repeated"
            total = max(1, sum(len(w) for w in ws))
            for n, cap in ((5, r["max_dup5gram"]), (10, r["max_dup10gram"])):
                if dup_ngram_chars(ws, n) / total > cap * rep:
                    return "repeated passages"
        return ""

    def _code(self, text: str) -> str:
        r = self.rules
        lines = text.split("\n")
        if len(text) < 100:
            return "too short"
        if max(len(l) for l in lines) > r["code_max_line"]:
            return "minified or data blob"
        if sum(c.isalpha() for c in text) / len(text) < r["code_min_alpha"]:
            return "too little real code"
        if AUTOGEN.search(text[:600]):
            return "auto-generated"
        return ""


# ---------------------------------------------------------------- decontamination

# Benchmarks people usually report. Each entry: (hub id, config, split, fields joined to form the test text).
BENCHMARKS = {
    "gsm8k": ("openai/gsm8k", "main", "test", ["question"]),
    "math500": ("HuggingFaceH4/MATH-500", "default", "test", ["problem"]),
    "humaneval": ("openai/openai_humaneval", "openai_humaneval", "test", ["prompt"]),
    "mbpp": ("google-research-datasets/mbpp", "full", "test", ["text"]),
    "mmlu": ("cais/mmlu", "all", "test", ["question"]),
    "arc_challenge": ("allenai/ai2_arc", "ARC-Challenge", "test", ["question"]),
    "arc_easy": ("allenai/ai2_arc", "ARC-Easy", "test", ["question"]),
    "hellaswag": ("Rowan/hellaswag", "default", "validation", ["ctx"]),
    "openbookqa": ("allenai/openbookqa", "main", "test", ["question_stem"]),
    "winogrande": ("allenai/winogrande", "winogrande_xl", "validation", ["sentence"]),
    "truthfulqa": ("truthfulqa/truthful_qa", "generation", "validation", ["question"]),
    "commonsenseqa": ("tau/commonsense_qa", "default", "validation", ["question"]),
}


class Decontaminator:
    """Drop a document when it shares any long n-gram with a held-out test item.

    13-grams is the GPT-3 / OLMo convention. Short test items (under n words) are matched as a whole normalized string.
    """

    def __init__(self, n: int = 13):
        self.n = n
        self.grams: dict[str, str] = {}
        self.short: dict[str, str] = {}
        self.counts: Counter = Counter()

    def add(self, name: str, texts: list[str]):
        for t in texts:
            ws = norm(t).split()
            if len(ws) < self.n:
                if len(ws) >= 6:
                    self.short[" ".join(ws)] = name
                continue
            for g in ngrams(ws, self.n):
                self.grams.setdefault(g, name)

    def __len__(self):
        return len(self.grams) + len(self.short)

    def hit(self, text: str) -> str:
        if not self:
            return ""
        nt = norm(text)
        ws = nt.split()
        for g in ngrams(ws, self.n):
            b = self.grams.get(g)
            if b:
                self.counts[b] += 1
                return b
        if self.short:
            for s, b in self.short.items():
                if s in nt:
                    self.counts[b] += 1
                    return b
        return ""


_BUNDLED: dict | None = None


def bundled_benchmarks() -> dict:
    """The test sets shipped with the package (data/benchmarks.json.gz, built by make_benchmarks.py)."""
    global _BUNDLED
    if _BUNDLED is None:
        import gzip
        import json
        import os
        path = os.path.join(os.path.dirname(__file__), "data", "benchmarks.json.gz")
        try:
            with gzip.open(path, "rt", encoding="utf-8") as f:
                _BUNDLED = json.load(f)
        except OSError:
            _BUNDLED = {}
    return _BUNDLED


def benchmark_texts(key: str, token: str = "") -> list[str]:
    """Bundled copy first; the dataset viewer only when a set is missing from the bundle."""
    return bundled_benchmarks().get(key) or load_benchmark_texts(key, token)


def load_benchmark_texts(key: str, token: str = "", limit: int = 20000) -> list[str]:
    """Test items through the dataset viewer API (no download). Test sets are small, so this is enough."""
    from .sources import _get
    ds, cfg, split, fields = BENCHMARKS[key]
    out, off = [], 0
    while off < limit:
        got = _get("rows", {"dataset": ds, "config": cfg, "split": split, "offset": off, "length": 100}, token)
        rows = got.get("rows", [])
        for r in rows:
            row = r.get("row") or {}
            t = " ".join(str(row.get(f, "")) for f in fields).strip()
            if t:
                out.append(t)
        if len(rows) < 100:
            break
        off += 100
    return out