#!/usr/bin/env python3
"""Extract clean Russian sentences from an official Wikisource XML dump."""

from __future__ import annotations

import argparse
import bz2
import heapq
import hashlib
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from xml.etree import ElementTree


SENTENCE_BOUNDARY = re.compile(r"(?<=[.!?…])\s+")
WHITESPACE = re.compile(r"\s+")
RUSSIAN_LETTER = re.compile(r"[А-Яа-яЁё]")
ANY_LETTER = re.compile(r"[^\W\d_]", re.UNICODE)
BAD_MARKUP = ("{{", "}}", "[[", "]]", "http://", "https://", "<math", "</", "{|", "|}")


def local_name(tag: str) -> str:
    return tag.rsplit("}", 1)[-1]


def descendant_text(element, name: str) -> str:
    for child in element.iter():
        if local_name(child.tag) == name:
            return child.text or ""
    return ""


def clean_page_text(source: str) -> str:
    import mwparserfromhell

    source = re.sub(r"<ref\b[^>]*>.*?</ref>|<ref\b[^>]*/>", " ", source, flags=re.IGNORECASE | re.DOTALL)
    source = re.sub(r"\{\|.*?\|\}", " ", source, flags=re.DOTALL)
    plain = mwparserfromhell.parse(source).strip_code(normalize=True, collapse=True)
    return WHITESPACE.sub(" ", plain).strip()


def acceptable(sentence: str) -> bool:
    if not 40 <= len(sentence) <= 240 or any(marker in sentence for marker in BAD_MARKUP):
        return False
    if sentence[-1:] not in ".!?…":
        return False
    words = sentence.split()
    if not 6 <= len(words) <= 45:
        return False
    letters = ANY_LETTER.findall(sentence)
    if not letters:
        return False
    russian = RUSSIAN_LETTER.findall(sentence)
    return len(russian) / len(letters) >= 0.75


def extract(args: argparse.Namespace) -> None:
    temporary = args.output.with_suffix(args.output.suffix + ".part")
    if args.output.exists() or args.metadata.exists() or temporary.exists():
        raise FileExistsError("Refusing to overwrite output, metadata, or partial output")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    seen: set[bytes] = set()
    sample_heap: list[tuple[int, bytes, str]] = []
    pages = 0
    accepted = 0
    candidates = 0
    with bz2.open(args.dump, "rb") as source, temporary.open("w", encoding="utf-8", newline="\n") as output:
        for _, element in ElementTree.iterparse(source, events=("end",)):
            if local_name(element.tag) != "page":
                continue
            pages += 1
            namespace = descendant_text(element, "ns")
            title = descendant_text(element, "title")
            wiki_text = descendant_text(element, "text")
            if namespace == "0" and wiki_text and not title.startswith(("Категория:", "Шаблон:", "Справка:")):
                for sentence in SENTENCE_BOUNDARY.split(clean_page_text(wiki_text)):
                    # Leading dashes carry dialogue/prosody information and must survive extraction.
                    sentence = sentence.strip(" \t\r\n")
                    if not acceptable(sentence):
                        continue
                    digest = hashlib.blake2b(sentence.encode("utf-8"), digest_size=16).digest()
                    candidates += 1
                    if args.sampling == "first":
                        if digest in seen:
                            continue
                        seen.add(digest)
                        output.write(sentence + "\n")
                        accepted += 1
                        if accepted % 10_000 == 0:
                            print(f"sentences={accepted} pages={pages}", flush=True)
                        if accepted >= args.max_sentences:
                            break
                    else:
                        score = int.from_bytes(digest, "big")
                        if digest in seen:
                            continue
                        if len(sample_heap) < args.max_sentences:
                            heapq.heappush(sample_heap, (-score, digest, sentence)); seen.add(digest)
                        elif score < -sample_heap[0][0]:
                            _, removed_digest, _ = heapq.heapreplace(sample_heap, (-score, digest, sentence))
                            seen.remove(removed_digest); seen.add(digest)
            element.clear()
            if args.sampling == "first" and accepted >= args.max_sentences:
                break
            if pages % 10_000 == 0:
                print(f"pages={pages} candidates={candidates} retained={len(sample_heap)}", flush=True)
        if args.sampling == "hash":
            selected = sorted((-negative_score, sentence) for negative_score, _, sentence in sample_heap)
            for _, sentence in selected:
                output.write(sentence + "\n")
            accepted = len(selected)
    temporary.replace(args.output)
    metadata = {
        "source_url": args.source_url,
        "dump_file": args.dump.name,
        "dump_bytes": args.dump.stat().st_size,
        "pages_scanned": pages,
        "candidate_sentences": candidates,
        "sentences": accepted,
        "output_bytes": args.output.stat().st_size,
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "sampling": args.sampling,
        "filters": {"min_chars": 40, "max_chars": 240, "min_words": 6, "max_words": 45, "min_russian_letter_ratio": 0.75},
    }
    args.metadata.write_text(json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dump", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--metadata", type=Path, required=True)
    parser.add_argument("--max-sentences", type=int, default=200_000)
    parser.add_argument(
        "--sampling", choices=("first", "hash"), default="hash",
        help="hash scans the complete dump and keeps a deterministic min-hash sample",
    )
    parser.add_argument(
        "--source-url",
        default="https://dumps.wikimedia.org/ruwikisource/latest/ruwikisource-latest-pages-articles-multistream.xml.bz2",
    )
    return parser.parse_args()


if __name__ == "__main__":
    sys.exit(extract(parse_args()) or 0)