"""Daten vorbereiten: echte Fehlerpaare laden, saubere Sätze sammeln, Tokenizer trainieren.

Kaggle-CPU-Session:  python prepare.py --out /kaggle/working/data
Lokaler Mini-Lauf:   python prepare.py --out data --max-sents 20000 --vocab-size 8000
"""
import argparse
import io
import json
import os
import re
import socket
import tarfile
import time
import urllib.request
import zipfile
from collections import Counter

TOKEN_RE = re.compile(r"\d+(?:[.,]\d+)+|\w+(?:[-'’]\w+)*|[^\w\s]")
SENT_SPLIT = re.compile(r"(?<=[.!?])\s+(?=[A-ZÄÖÜ„\"])")
ALLOWED = re.compile(r"[\w\s.,;:!?\-–'’\"„“‚‘()%/&]+")
BOYD_URL = "https://github.com/adrianeboyd/boyd-wnut2018/releases/download/wnut2018/data.tar.gz"
OWN_URL = "https://huggingface.co/datasets/Bayernator/german-spelling-data"
BENCH_URL = "https://huggingface.co/datasets/Bayernator/german-ortho/resolve/main/rechtschreibung_hard_benchmark.zip"
ABBREV = set("bzw usw ca vgl z B d h u a o v n Chr Nr Dr Prof St etc geb gest sog evtl ggf inkl Jh Mio Mrd bspw Hrsg Bd S f ff Abb Tab".split())
REAL_FILES = ["fm-train.src", "fm-train.trg", "fm-dev.src", "fm-dev.trg", "fm-dev.m2",
              "fm-test.src", "fm-test.trg", "fm-test.m2", "wiki.1M.src", "wiki.1M.trg"]


_getaddrinfo = socket.getaddrinfo
socket.getaddrinfo = lambda *a, **k: sorted(_getaddrinfo(*a, **k), key=lambda x: x[0] != socket.AF_INET)  # IPv4 zuerst: kaputtes IPv6 lässt urllib sonst ewig hängen


def get(url):
    return urllib.request.urlopen(url, timeout=120).read()


def tokenize(text):
    # ponytail: Regex-Tokenizer, Abkürzungen ("z. B.", "bzw.") werden gesplittet; reicht, weil Quelle und Ziel gleich tokenisiert sind
    return " ".join(TOKEN_RE.findall(text))


def detokenize(text):
    text = re.sub(r" ([.,;:!?)\]}“%])", r"\1", text)
    text = re.sub(r"([(\[{„]) ", r"\1", text)
    text = re.sub(r"(\w) - (und|oder|bzw)\b", r"\1- \2", text)  # Ergänzungsstrich: "Manga- und Animeserien"
    out, open_q = [], {'"': False, "'": False}
    for tok in text.split():
        if tok in open_q:  # gerade Anführungszeichen paarweise: öffnend klebt am nächsten Wort, schließend am vorherigen
            open_q[tok] = not open_q[tok]
            if open_q[tok]:
                out.append(tok + "\0")  # Marker: kein Leerzeichen danach
            elif out:
                out[-1] += tok
            else:
                out.append(tok)
        else:
            out.append(tok)
    return " ".join(out).replace("\0 ", "").replace("\0", "")


def split_sentences(text):
    for line in text.split("\n"):
        yield from SENT_SPLIT.split(line.strip())


def good(s):
    return (20 <= len(s) <= 300 and s[0].isupper() and s[-1] in ".!?" and ALLOWED.fullmatch(s)
            and sum(c.isdigit() for c in s) < 0.1 * len(s))


def fetch_real(out):
    os.makedirs(f"{out}/real", exist_ok=True)
    if all(os.path.exists(f"{out}/real/{f}") for f in REAL_FILES):
        return
    print("Lade Falko-MERLIN + Wikipedia-Edits (Boyd 2018, CC BY-SA 4.0) ...")
    blob = get(BOYD_URL)
    with tarfile.open(fileobj=io.BytesIO(blob)) as tar:
        for f in REAL_FILES:
            with open(f"{out}/real/{f}", "wb") as fh:
                fh.write(tar.extractfile(f"data/{f}").read())


def _pair_split(text):
    """Satztrennung für Absatzpaare, die auch klein geschriebene Satzanfänge (Fehler!) trennt."""
    out = []
    for piece in re.split(r"(?<=[.!?])\s+", text.strip()):
        last = re.split(r"[\s(]", out[-1])[-1][:-1] if out else ""
        if out and (last in ABBREV or last.isdigit() or len(last) <= 1):  # "v. Chr.", "3. Februar"
            out[-1] += " " + piece
        else:
            out.append(piece)
    return out


def sentence_pairs(inp, tgt):
    """Absatzpaar -> Satzpaare; der letzte Satz fällt weg, wenn der Absatz mitten im Text abgeschnitten ist."""
    ss, ts = _pair_split(inp), _pair_split(tgt)
    if len(ss) != len(ts):
        return []
    if ts[-1][-1:] not in ".!?\"“)":
        ss, ts = ss[:-1], ts[:-1]
    return [(tokenize(s), tokenize(t)) for s, t in zip(ss, ts) if len(t.split()) <= 80]


def fetch_own(out):
    """Eigene HF-Datasets: Hard-Benchmark (nur Evaluation) + Trainingspaare ohne Benchmark-Überschneidung."""
    if os.path.exists(f"{out}/real/own.src"):
        return
    print("Lade Bayernator/german-spelling-data + Hard-Benchmark ...")
    z = zipfile.ZipFile(io.BytesIO(get(BENCH_URL)))
    bench = z.read("rechtschreibung_hard_benchmark/benchmark_hard_5000.jsonl").decode()
    with open(f"{out}/real/bench.jsonl", "w", encoding="utf-8") as fh:
        fh.write(bench)
    rows = [json.loads(l) for l in bench.splitlines() if l.strip()]
    banned = {r["input"] for r in rows} | {r["target"] for r in rows}
    files = json.loads(get(OWN_URL.replace("/datasets/", "/api/datasets/") + "/tree/main/data"))
    kept = dropped = 0
    with open(f"{out}/real/own.src", "w", encoding="utf-8") as fs, open(f"{out}/real/own.trg", "w", encoding="utf-8") as ft:
        for f in files:
            for line in get(f"{OWN_URL}/resolve/main/{f['path']}").decode().splitlines():
                r = json.loads(line) if line.strip() else None
                if not r or r["input"] in banned or r["target"] in banned:
                    dropped += bool(r)
                    continue
                for s, t in sentence_pairs(r["input"], r["target"]):
                    fs.write(s + "\n")
                    ft.write(t + "\n")
                    kept += 1
    print(f"{kept:,} eigene Satzpaare, {dropped:,} Absätze wegen Benchmark-Überschneidung entfernt")


def collect_clean(out, max_sents, fineweb_share, max_hours=8.0):
    import numpy as np
    from datasets import load_dataset
    sources = [("wikimedia/wikipedia", "20231101.de", int(max_sents * (1 - fineweb_share))),
               ("epfml/FineWeb2-HQ", "deu_Latn", max_sents)]  # FineWeb füllt bis max_sents auf
    seen, offsets, words, pos = set(), [], Counter(), 0
    deadline = time.time() + max_hours * 3600  # Kaggle-Session: 12 h, Tokenizer-Training braucht danach noch Zeit
    with open(f"{out}/clean.txt", "wb") as fh:
        for name, config, quota in sources:
            for doc in load_dataset(name, config, split="train", streaming=True):
                for s in split_sentences(doc["text"]):
                    if not good(s):
                        continue
                    t = tokenize(s)
                    n = t.count(" ") + 1
                    if not 5 <= n <= 80 or hash(t) in seen:
                        continue
                    seen.add(hash(t))
                    if len(offsets) < 5_000_000:  # Wortschatz aus den ersten 5 Mio. Sätzen reicht
                        words.update(w for w in t.split() if w.isalpha())
                    b = (t + "\n").encode()
                    offsets.append(pos)
                    fh.write(b)
                    pos += len(b)
                    if len(offsets) % 1_000_000 == 0:
                        print(f"{len(offsets):,} Sätze")
                if len(offsets) >= quota or time.time() > deadline:
                    break
    offsets.append(pos)  # Ende des letzten Satzes
    np.save(f"{out}/clean.idx.npy", np.array(offsets, dtype=np.uint64))
    min_count = 5 if len(offsets) > 1_000_000 else 2
    with open(f"{out}/vocab.txt", "w") as fh:
        fh.writelines(f"{w}\n" for w, c in words.most_common(100_000) if c >= min_count)
    print(f"{len(offsets) - 1:,} saubere Sätze gespeichert")


def train_tokenizer(out, vocab_size):
    import sentencepiece as spm
    spm.SentencePieceTrainer.train(
        input=f"{out}/clean.txt", model_prefix=f"{out}/spm", vocab_size=vocab_size,
        model_type="unigram", character_coverage=0.9995, byte_fallback=True, split_digits=True,
        normalization_rule_name="identity", input_sentence_size=5_000_000, shuffle_input_sentence=True,
        pad_id=0, unk_id=1, bos_id=2, eos_id=3, num_threads=os.cpu_count(), minloglevel=2)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data")
    ap.add_argument("--max-sents", type=int, default=30_000_000)
    ap.add_argument("--fineweb-share", type=float, default=0.5)
    ap.add_argument("--vocab-size", type=int, default=32000)
    ap.add_argument("--max-hours", type=float, default=8.0, help="Zeitbudget fürs Sammeln sauberer Sätze")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    fetch_real(a.out)
    fetch_own(a.out)
    collect_clean(a.out, a.max_sents, a.fineweb_share, a.max_hours)
    train_tokenizer(a.out, a.vocab_size)
    print("Fertig:", a.out)