"""Phase 2.3 — clean every raw bucket → data/cpt_clean/<bucket>.jsonl

Pipeline order (plan §2.3), applied per doc, then a global cross-bucket dedup:
  1. ftfy unicode fix (NFC)
  2. fastText language id — drop if top lang != ru or p(ru) < 0.65
  3. mixed-script rejection — drop docs with homoglyph-contaminated words
  4. boilerplate/length — strip HTML, collapse whitespace, drop <200 or >50k chars,
     drop docs with >30% non-letter chars
  5. MinHash-LSH near-dup removal across ALL buckets (the biggest quality lever)
  6. light PII scrub (phones/emails) — logs bucket has real data

Usage:  python corpus_clean.py            # all buckets present in data/cpt_raw/
        python corpus_clean.py wiki synth # specific buckets
"""
import sys
import re
import json
import html

import ftfy
from datasketch import MinHash, MinHashLSH

from corpus_common import RAW, CLEAN, has_mixed_script, lang_prob

MIN_CHARS, MAX_CHARS = 200, 50_000
NON_LETTER_MAX = 0.30
LID_MIN = 0.65
LSH_THRESHOLD = 0.8
NUM_PERM = 128

_TAG = re.compile(r"<[^>]+>")
_WS = re.compile(r"[ \t ]+")
_NL = re.compile(r"\n{3,}")
_LETTER = re.compile(r"[^\W\d_]", re.UNICODE)
_EMAIL = re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+")
_PHONE = re.compile(r"(?<!\d)(?:\+7|8)[\s(-]*\d{3}[\s)-]*\d{3}[\s-]*\d{2}[\s-]*\d{2}(?!\d)")


def normalize(text: str) -> str:
    text = ftfy.fix_text(text)
    text = html.unescape(text)
    text = _TAG.sub(" ", text)
    text = _WS.sub(" ", text)
    text = _NL.sub("\n\n", text)
    return text.strip()


def non_letter_ratio(text: str) -> float:
    letters = len(_LETTER.findall(text))
    return 1.0 - (letters / len(text)) if text else 1.0


def scrub_pii(text: str) -> str:
    text = _EMAIL.sub("<email>", text)
    text = _PHONE.sub("<phone>", text)
    return text


def doc_minhash(text: str) -> MinHash:
    m = MinHash(num_perm=NUM_PERM)
    for tok in text.lower().split():
        m.update(tok.encode("utf-8"))
    return m


def clean_bucket(bucket: str, lsh: MinHashLSH, seen: dict) -> int:
    src = RAW / bucket / "shard-000.jsonl"
    if not src.exists():
        print(f"[{bucket}] no raw shard at {src}; skip")
        return 0
    out = CLEAN / f"{bucket}.jsonl"
    kept = dropped = 0
    with open(src, encoding="utf-8") as fin, open(out, "w", encoding="utf-8") as fout:
        for i, line in enumerate(fin):
            try:
                text = json.loads(line)["text"]
            except (json.JSONDecodeError, KeyError):
                dropped += 1
                continue
            text = normalize(text)
            # 4. length / boilerplate
            if not (MIN_CHARS <= len(text) <= MAX_CHARS) or non_letter_ratio(text) > NON_LETTER_MAX:
                dropped += 1
                continue
            # 3. mixed script
            if has_mixed_script(text):
                dropped += 1
                continue
            # 2. language id
            try:
                lang, p = lang_prob(text)
                if lang != "ru" or p < LID_MIN:
                    dropped += 1
                    continue
            except FileNotFoundError:
                raise SystemExit("lid.176.bin missing — see corpus_common.LID_PATH")
            # 5. dedup
            m = doc_minhash(text)
            if lsh.query(m):
                dropped += 1
                continue
            key = f"{bucket}:{i}"
            lsh.insert(key, m)
            seen[key] = True
            # 6. pii
            text = scrub_pii(text)
            fout.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
            kept += 1
            if (kept + dropped) % 5000 == 0:
                print(f"[{bucket}] kept {kept} dropped {dropped}")
    print(f"[{bucket}] DONE kept {kept} dropped {dropped} -> {out}")
    return kept


def main():
    buckets = sys.argv[1:] or [p.name for p in sorted(RAW.iterdir()) if p.is_dir()]
    # one shared LSH across all buckets = global dedup (plan §2.3 rule 5)
    lsh = MinHashLSH(threshold=LSH_THRESHOLD, num_perm=NUM_PERM)
    seen = {}
    total = 0
    for b in buckets:
        total += clean_bucket(b, lsh, seen)
    print(f"[clean] TOTAL kept {total} docs across {len(buckets)} buckets")


if __name__ == "__main__":
    main()