"""Shared paths + cleaning helpers for the RU-CPT corpus pipeline.

The mixed-script / Cyrillic-cleanliness rule lives here ONCE and is reused by
corpus_clean.py, corpus_synth.py, and eval_cyrillic.py. See
docs/plans/2026-07-11-ru-cpt-base.md (§2.3 rule 3) in the Maven repo.

Run `python corpus_common.py` for a self-check of the mixed-script detector.
"""
import os
import re
from pathlib import Path

# ── env / cache (mirror train_rocm.py) ──────────────────────────────────────
os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")
os.makedirs(os.environ["TMPDIR"], exist_ok=True)

# ── corpus dir layout ───────────────────────────────────────────────────────
DATA = Path(__file__).resolve().parent / "data"
RAW = DATA / "cpt_raw"        # corpus_fetch.py / corpus_synth.py output
CLEAN = DATA / "cpt_clean"    # corpus_clean.py output (JSONL {"text":...})
PACKED = DATA / "cpt_packed"  # corpus_pack.py output (HF Arrow, 2048-tok blocks)
for d in (RAW, CLEAN, PACKED):
    d.mkdir(parents=True, exist_ok=True)

# ── the corpus mix (plan §2.1) — bucket -> target token share ───────────────
TARGET_TOKENS = 300_000_000
BUCKET_SHARE = {          # sums to 1.0; synth capped at 0.10 (hard rule)
    "culturax": 0.60,
    "wiki": 0.17,
    "books": 0.11,
    "synth": 0.10,
    "logs": 0.02,
}
SYNTH_CAP = 0.10

# ── Cyrillic cleanliness (plan §2.3 rule 3) ─────────────────────────────────
_WORD = re.compile(r"[^\s]+")
_CYR = re.compile(r"[а-яёА-ЯЁ]")
_LAT = re.compile(r"[a-zA-Z]")


def is_mixed_script_word(word: str) -> bool:
    """True if a single token mixes Cyrillic and Latin letters (homoglyph swap)."""
    return bool(_CYR.search(word) and _LAT.search(word))


def has_mixed_script(text: str) -> bool:
    """True if ANY word in the text is homoglyph-contaminated."""
    return any(is_mixed_script_word(w) for w in _WORD.findall(text))


def cyrillic_ratio(text: str) -> float:
    """Fraction of letters that are Cyrillic (0..1). Empty/letterless -> 0."""
    cyr = len(_CYR.findall(text))
    lat = len(_LAT.findall(text))
    tot = cyr + lat
    return cyr / tot if tot else 0.0


# ── fastText language id (lazy singleton; download lid.176.bin, plan §2.3) ──
_LID = None
LID_PATH = os.environ.get("FASTTEXT_LID", str(DATA / "lid.176.bin"))


def lang_prob(text: str):
    """Return (lang_code, prob) for the dominant language, or (None, 0.0)."""
    global _LID
    if _LID is None:
        import fasttext
        if not Path(LID_PATH).exists():
            raise FileNotFoundError(
                f"fastText lid.176.bin not found at {LID_PATH}. Download from "
                "https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin "
                "or set FASTTEXT_LID env var."
            )
        # numpy 2.x compat: fasttext uses np.array(..., copy=False) which
        # raises on numpy >=2 when a copy is needed. patch its np reference.
        _orig_np_array = fasttext.FastText.np.array
        def _patched_array(*args, **kwargs):
            kwargs.pop("copy", None)
            return _orig_np_array(*args, **kwargs)
        fasttext.FastText.np.array = _patched_array
        _LID = fasttext.load_model(LID_PATH)
    labels, probs = _LID.predict(text.replace("\n", " ")[:2000])
    return labels[0].replace("__label__", ""), float(probs[0])


if __name__ == "__main__":
    # self-check: the homoglyph detector is the whole point — assert it works.
    assert is_mixed_script_word("privет")          # latin+cyrillic in one word
    assert is_mixed_script_word("Марkет")           # cyrillic 'Мар' + latin 'k'
    assert not is_mixed_script_word("привет")       # pure cyrillic
    assert not is_mixed_script_word("hello")        # pure latin
    assert not is_mixed_script_word("http://x.ru")  # pure latin url — keep
    assert has_mixed_script("это Марkет здесь")
    assert not has_mixed_script("это чистый русский текст with english words")
    assert abs(cyrillic_ratio("привет hello") - 0.5) < 0.05
    print("corpus_common self-check OK ->", DATA)