#!/usr/bin/env python3
"""Массовая заливка русской Википедии в корпус сравнения.

Зачем именно она: студенты копируют из Википедии чаще, чем из научных статей, а
в корпусе её нет вовсе. При этом она единственный русскоязычный источник такого
объёма, который отдаётся без блокировок — КиберЛенинка режет выкачку после ~130
запросов, eLIBRARY требует договора. Дамп `ruwiki-latest-pages-articles.xml.bz2`
(5.6 ГБ, ~2 млн статей) скачивается свободно.

Дамп не сохраняется на диск: читается потоком и распаковывается на лету —
на app-хосте всего 22 ГБ свободно.

Что отбрасывается: перенаправления, служебные пространства имён (оставляем
только статьи) и короткие тексты — для детекции нужен объём, а не заготовки.

Запуск (в контейнере worker-indexer):
    cd /home/user/anti-plagiarism
    C="docker compose -f docker-compose.prod.yml exec -T worker-indexer python -"
    $C --limit 100 < scripts/ops/bulk_ingest_wikipedia_ru.py       # проба
    $C --limit 500000 < scripts/ops/bulk_ingest_wikipedia_ru.py    # объём

Идемпотентно: ext_id вида `wikipedia_ru:<pageid>`, повторная заливка пропускает
уже существующие статьи (ON CONFLICT DO NOTHING).
"""

import argparse
import bz2
import io
import re
import time

DUMP_URL = "https://dumps.wikimedia.org/ruwiki/latest/ruwiki-latest-pages-articles.xml.bz2"
PAGE_RE = re.compile(r"<page>(.*?)</page>", re.DOTALL)
TITLE_RE = re.compile(r"<title>(.*?)</title>", re.DOTALL)
ID_RE = re.compile(r"<id>(\d+)</id>")
NS_RE = re.compile(r"<ns>(\d+)</ns>")
TEXT_RE = re.compile(r'<text[^>]*>(.*?)</text>', re.DOTALL)
REDIRECT_RE = re.compile(r"<redirect ")

# Разметку чистим регулярками: для отпечатков нужен связный текст, а не точное
# восстановление вёрстки — полноценный парсер вики-разметки тут неоправдан
CLEAN_RULES = [
    (re.compile(r"\{\{[^{}]*\}\}"), " "),          # шаблоны (в несколько проходов)
    (re.compile(r"\[\[[^\]|]*\|"), ""),             # [[ссылка|текст → текст
    (re.compile(r"\[\[|\]\]"), ""),                 # остатки скобок
    (re.compile(r"<ref[^>]*>.*?</ref>", re.DOTALL), " "),
    (re.compile(r"<[^>]+>"), " "),                  # html-теги
    (re.compile(r"^[*#:;|!].*$", re.MULTILINE), " "),  # списки и таблицы
    (re.compile(r"^=+.*?=+$", re.MULTILINE), " "),     # == заголовки разделов ==
    (re.compile(r"'{2,}"), ""),                     # ''курсив''
    (re.compile(r"&[a-z]+;"), " "),
    (re.compile(r"[ \t]+"), " "),
    (re.compile(r"\n{2,}"), "\n"),
]


def clean_wikitext(raw: str) -> str:
    text = raw
    for _ in range(3):  # шаблоны бывают вложенными
        text = CLEAN_RULES[0][0].sub(CLEAN_RULES[0][1], text)
    for rx, repl in CLEAN_RULES[1:]:
        text = rx.sub(repl, text)
    return text.strip()


def iter_pages(min_chars: int):
    """Потоково читать дамп и отдавать статьи основного пространства имён."""
    import httpx

    decomp = bz2.BZ2Decompressor()
    buf = ""
    # Wikimedia отдаёт 403 без осмысленного User-Agent — по их правилам он должен
    # называть приложение и давать контакт
    headers = {"User-Agent": "AcademicHelper/1.0 (https://academic.jze9.ru; noreply@jze9.ru)"}
    with httpx.stream("GET", DUMP_URL, timeout=120, follow_redirects=True, headers=headers) as resp:
        resp.raise_for_status()
        for chunk in resp.iter_bytes(4 * 1024 * 1024):
            try:
                raw = decomp.decompress(chunk)
            except EOFError:
                break
            if not raw:
                continue
            buf += raw.decode("utf-8", errors="replace")

            while True:
                m = PAGE_RE.search(buf)
                if not m:
                    break
                page, buf = m.group(1), buf[m.end():]

                if REDIRECT_RE.search(page):
                    continue
                ns = NS_RE.search(page)
                if not ns or ns.group(1) != "0":      # только статьи
                    continue
                tm, im, xm = TITLE_RE.search(page), ID_RE.search(page), TEXT_RE.search(page)
                if not (tm and im and xm):
                    continue
                text = clean_wikitext(xm.group(1))
                if len(text) < min_chars:
                    continue
                yield im.group(1), tm.group(1), text

            if len(buf) > 20 * 1024 * 1024:            # страховка от разбухания
                buf = buf[-1024 * 1024:]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--limit", type=int, default=100, help="сколько статей залить")
    ap.add_argument("--batch", type=int, default=500, help="статей в транзакции")
    ap.add_argument("--min-chars", type=int, default=2000, help="минимальная длина текста")
    ap.add_argument("--fp-per-doc", type=int, default=2000, help="максимум отпечатков на статью")
    args = ap.parse_args()

    import psycopg2
    from app.algorithms.winnowing import winnow
    from app.config import settings

    conn = psycopg2.connect(host=settings.POSTGRES_HOST, port=settings.POSTGRES_PORT,
                            dbname=settings.POSTGRES_DB, user=settings.POSTGRES_USER,
                            password=settings.POSTGRES_PASSWORD)
    added = skipped = 0
    fp_total = 0
    t0 = time.time()
    batch: list[tuple[str, str, str]] = []

    def flush(batch):
        nonlocal added, skipped, fp_total
        if not batch:
            return
        with conn.cursor() as cur:
            cur.executemany(
                """INSERT INTO documents (ext_id, source, title, lang, abstract, url,
                                          authors, indexed_at)
                   VALUES (%s, 'wikipedia_ru', %s, 'ru', %s, %s, '[]'::json, now())
                   ON CONFLICT (ext_id) DO NOTHING""",
                [(f"wikipedia_ru:{pid}", title[:1000], text[:2000],
                  f"https://ru.wikipedia.org/?curid={pid}") for pid, title, text in batch],
            )
            cur.execute("SELECT id, ext_id FROM documents WHERE ext_id = ANY(%s)",
                        ([f"wikipedia_ru:{pid}" for pid, _, _ in batch],))
            id_by_ext = {e: i for i, e in cur.fetchall()}

            buf = io.StringIO()
            fresh = 0
            for pid, _, text in batch:
                doc_id = id_by_ext.get(f"wikipedia_ru:{pid}")
                if doc_id is None:
                    continue
                cur.execute("SELECT 1 FROM fingerprints WHERE doc_id=%s LIMIT 1", (doc_id,))
                if cur.fetchone():
                    skipped += 1
                    continue
                hashes = list(winnow(text))[: args.fp_per_doc]
                for pos, h in enumerate(hashes):
                    buf.write(f"{doc_id}\t{h}\t{pos}\n")
                fp_total += len(hashes)
                fresh += 1
            buf.seek(0)
            if fresh:
                cur.copy_from(buf, "fingerprints", columns=("doc_id", "hash_value", "position"))
            added += fresh
        conn.commit()

    for pid, title, text in iter_pages(args.min_chars):
        batch.append((pid, title, text))
        if len(batch) >= args.batch:
            flush(batch)
            batch = []
            el = time.time() - t0
            print(f"  залито {added} · пропущено {skipped} · отпечатков {fp_total:,} · "
                  f"{added / max(el, 1):.1f} ст/с", flush=True)
            if added >= args.limit:
                break
    flush(batch)
    conn.close()

    el = time.time() - t0
    print(f"\nГотово за {el / 60:.1f} мин: залито {added}, пропущено {skipped}, "
          f"отпечатков {fp_total:,}")


if __name__ == "__main__":
    main()