"""Streaming adapters for each pretraining source. Each yields cleaned text documents."""
from __future__ import annotations

import io
import json
import random
import re
from pathlib import Path

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
import yaml
from huggingface_hub import HfApi, hf_hub_download

ROOT = Path(__file__).resolve().parent.parent
RAW = ROOT / "data" / "raw"
_BLANKS = re.compile(r"\n{4,}")


def load_cfg(path=None):
    return yaml.safe_load(open(path or ROOT / "configs" / "data.yaml"))


def _clean(t: str) -> str:
    t = t.replace("\r\n", "\n").replace("\x00", "")
    return _BLANKS.sub("\n\n\n", t).strip()


class Source:
    def __init__(self, name: str, cfg: dict, glob: dict):
        self.name, self.cfg, self.glob = name, cfg, glob
        self.repo = cfg["repo"]
        self.min_chars = cfg.get("min_chars", glob.get("min_chars", 200))
        self.quota = int(glob["total_tokens"] * cfg["weight"])
        self.per_shard_cap = cfg.get("per_shard_cap_tokens")
        self.api = HfApi()

    # -- file selection (deterministic shuffle so a partial run is still a uniform sample) --
    def list_shards(self) -> list[str]:
        files = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=self.cfg["prefix"].rstrip("/"), recursive=True)
                 if hasattr(f, "size") and f.path.endswith((".parquet", ".jsonl.zst"))]
        files.sort()
        random.Random(self.glob["seed"]).shuffle(files)
        return files

    def download(self, path: str) -> Path:
        return Path(hf_hub_download(self.repo, path, repo_type="dataset", local_dir=RAW / self.name))

    def docs(self, local: Path):
        raise NotImplementedError

    def _emit(self, texts):
        for t in texts:
            if t is None:
                continue
            t = _clean(t)
            if len(t) >= self.min_chars:
                yield t


class ParquetSource(Source):
    text_col = "text"
    extra_cols: list[str] = []

    def mask(self, tbl: pa.Table):
        return None

    def format(self, tbl: pa.Table):
        return tbl[self.text_col].to_pylist()

    def docs(self, local: Path):
        pf = pq.ParquetFile(local)
        cols = [self.text_col] + self.extra_cols
        for batch in pf.iter_batches(batch_size=4096, columns=cols):
            tbl = pa.Table.from_batches([batch])
            m = self.mask(tbl)
            if m is not None:
                tbl = tbl.filter(m)
            yield from self._emit(self.format(tbl))


class FineWebEdu(ParquetSource):
    extra_cols = ["metadata"]

    def mask(self, tbl):
        score = pc.struct_field(tbl["metadata"], "score")
        return pc.greater_equal(score, self.cfg["min_edu_score"])


class Cosmopedia(ParquetSource):
    pass


class FineMath(ParquetSource):
    extra_cols = ["score", "language", "language_score"]

    def mask(self, tbl):
        return pc.and_(pc.equal(tbl["language"], "en"), pc.greater_equal(tbl["language_score"], 0.9))


class Wikipedia(ParquetSource):
    extra_cols = ["title"]

    def format(self, tbl):
        out = []
        for title, text in zip(tbl["title"].to_pylist(), tbl["text"].to_pylist()):
            if title.startswith(("List of", "Lists of", "Index of")) or len(text) < self.cfg.get("min_chars", 2000):
                continue
            out.append(f"# {title}\n\n{text}")
        return out


class DCLM(ParquetSource):
    extra_cols = ["fasttext_score", "language", "language_score"]

    def list_shards(self):
        base = "filtered/OH_eli5_vs_rw_v2_bigram_200k_train/fasttext_openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train/processed_data"
        rng = random.Random(self.glob["seed"])
        out = [f"{base}/global-shard_{g:02d}_of_10/local-shard_{l}_of_10/shard_{n:08d}_processed.parquet"
               for g in range(1, 11) for l in range(10) for n in range(0, 260)]
        rng.shuffle(out)
        return out[:self.cfg.get("max_shards", 200)]  # far more than needed; builder stops at quota

    def mask(self, tbl):
        m = pc.and_(pc.greater_equal(tbl["fasttext_score"], self.cfg["min_fasttext"]),
                    pc.greater_equal(tbl["language_score"], self.cfg["min_lang_score"]))
        return pc.and_(m, pc.equal(tbl["language"], "en"))


class StackExchange(ParquetSource):
    text_col = "ThreadText"
    extra_cols = ["Score", "AnswerCount"]

    def list_shards(self):
        out = []
        for site in self.cfg["sites"]:
            d = site if "." in site else f"{site}.stackexchange.com"
            if site == "stackoverflow":
                d = "stackoverflow.com"
            try:
                fs = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=d) if f.path.endswith(".parquet")]
            except Exception:
                continue
            off, n = self.cfg.get("shard_offset", 0), self.cfg.get("shards_per_site", 1)
            out.extend(sorted(fs)[off:off + n])  # few shards per site; per_shard_cap keeps the mix balanced
        random.Random(self.glob["seed"]).shuffle(out)
        return out

    def mask(self, tbl):
        return pc.and_(pc.greater_equal(tbl["Score"], self.cfg["min_score"]), pc.greater_equal(tbl["AnswerCount"], 1))


class CodePython(Source):
    def docs(self, local: Path):
        import zstandard
        with open(local, "rb") as fh, zstandard.ZstdDecompressor().stream_reader(fh) as r:
            for line in io.TextIOWrapper(r, encoding="utf-8"):
                o = json.loads(line)
                if o.get("score", 0) < self.cfg["min_score"] or o.get("length_bytes", 0) < self.cfg["min_bytes"]:
                    continue
                yield from self._emit([o["text"]])


REGISTRY = {"fineweb_edu": FineWebEdu, "dclm": DCLM, "cosmopedia": Cosmopedia, "finemath": FineMath,
            "wikipedia": Wikipedia, "stackexchange": StackExchange, "code_python": CodePython}


def get_source(name: str, cfg: dict | None = None) -> Source:
    cfg = cfg or load_cfg()
    return REGISTRY[name](name, cfg["sources"][name], cfg)