#!/usr/bin/env python3
"""
scrape.py — Download and process all Necyklopedie articles.

Three phases:
  1. Enumerate all article titles via MediaWiki API
  2. Bulk export via Special:Export (XML dump with full wikitext)
  3. Parse wikitext to plaintext and save as JSONL
"""

import json
import os
import re
import sys
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

API_URL = "https://necyklopedie.org/w/api.php"
EXPORT_URL = "https://necyklopedie.org/wiki/Speci%C3%A1ln%C3%AD:Exportovat_str%C3%A1nky"
OUTPUT_DIR = "data"
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "articles.jsonl")
TITLES_CACHE = os.path.join(OUTPUT_DIR, "titles.json")
DUMP_FILE = os.path.join(OUTPUT_DIR, "dump.xml")
MW_NS = {"mw": "http://www.mediawiki.org/xml/export-0.11/"}
UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"


def api_get(params, timeout=30, max_bytes=None):
    """One MediaWiki API call; `max_bytes` caps the reply (ValueError past it)."""
    params["format"] = "json"
    url = f"{API_URL}?{urllib.parse.urlencode(params)}"
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        payload = resp.read(-1 if max_bytes is None else max_bytes + 1)
    if max_bytes is not None and len(payload) > max_bytes:
        raise ValueError(f"API reply too large (> {max_bytes} B)")
    return json.loads(payload.decode("utf-8"))


def fetch_pages(titles, timeout=30, max_bytes=None):
    """Current text of these pages: {pageid: {pageid, title, text}}.

    Redirects, missing pages and stubs under 50 characters are left out. Uses
    the API, 50 titles a call: since 2026-10 necyklopedie.org answers a POST
    to Special:Export with 405.
    """
    results = {}
    for i in range(0, len(titles), 50):
        params = {
            "action": "query", "prop": "revisions", "rvprop": "content",
            "rvslots": "main", "formatversion": "2", "titles": "|".join(titles[i:i + 50]),
        }
        while True:
            data = api_get(params, timeout, max_bytes)
            for page in data.get("query", {}).get("pages", []):
                if page.get("redirect") or "revisions" not in page:
                    continue  # a page held back by a size limit comes after `continue`
                text = wikitext_to_plaintext(page["revisions"][0]["slots"]["main"]["content"])
                if len(text) >= 50:
                    results[page["pageid"]] = {
                        "pageid": page["pageid"], "title": page["title"], "text": text,
                    }
            if "continue" not in data:
                break
            params.update(data["continue"])
    return results


def get_all_titles():
    """Enumerate all article titles in namespace 0."""
    if os.path.exists(TITLES_CACHE):
        with open(TITLES_CACHE, "r", encoding="utf-8") as f:
            titles = json.load(f)
        print(f"Nalezeno {len(titles)} titulků v keši. Necyklopedie si pamatuje.", file=sys.stderr)
        return titles

    titles = []
    params = {"action": "query", "list": "allpages", "apnamespace": "0", "aplimit": "500"}

    while True:
        data = api_get(params)
        for page in data["query"]["allpages"]:
            titles.append({"pageid": page["pageid"], "title": page["title"]})
        print(f"\rProbírám se seznamem článků... {len(titles)}", end="", file=sys.stderr)
        if "continue" not in data:
            break
        params["apcontinue"] = data["continue"]["apcontinue"]
        time.sleep(0.5)

    print(f"\nNalezeno {len(titles)} článků. Tolik pravdy na jednom místě!", file=sys.stderr)
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    with open(TITLES_CACHE, "w", encoding="utf-8") as f:
        json.dump(titles, f, ensure_ascii=False)
    return titles


def bulk_export(titles):
    """Export all articles via Special:Export in one POST request."""
    if os.path.exists(DUMP_FILE):
        print(f"Používám keš ({DUMP_FILE}). Necyklopedie se stahuje jen jednou.", file=sys.stderr)
        return

    all_titles = "\n".join(t["title"] for t in titles)
    print(f"Stahuji {len(titles)} stránek přes Special:Export — tohle chvíli potrvá...", file=sys.stderr)

    data = urllib.parse.urlencode({
        "catname": "",
        "pages": all_titles,
        "curonly": "1",
        "wpDownload": "1",
    }).encode()
    req = urllib.request.Request(EXPORT_URL, data=data, headers={"User-Agent": UA})
    resp = urllib.request.urlopen(req, timeout=300)
    content = resp.read()

    with open(DUMP_FILE, "wb") as f:
        f.write(content)
    print(f"Staženo {len(content) / 1e6:.1f} MB čisté moudrosti do {DUMP_FILE}.", file=sys.stderr)


def wikitext_to_plaintext(wikitext):
    """Convert MediaWiki wikitext to clean plaintext."""
    text = wikitext

    # Remove categories, interwiki links
    text = re.sub(r"\[\[Kategorie:[^\]]*\]\]", "", text)
    text = re.sub(r"\[\[[a-z]{2,3}:[^\]]*\]\]", "", text)

    # Remove files/images: [[Soubor:...]] [[File:...]] [[Image:...]]
    text = re.sub(r"\[\[(?:Soubor|File|Image|Obrázek):[^\]]*\]\]", "", text, flags=re.IGNORECASE)

    # Remove templates {{...}} — handle nested braces
    for _ in range(5):  # multiple passes for nesting
        text = re.sub(r"\{\{[^{}]*\}\}", "", text)

    # Remove HTML tags
    text = re.sub(r"<ref[^>]*>.*?</ref>", "", text, flags=re.DOTALL)
    text = re.sub(r"<ref[^/]*/?>", "", text)
    text = re.sub(r"<gallery[^>]*>.*?</gallery>", "", text, flags=re.DOTALL)
    text = re.sub(r"<nowiki>.*?</nowiki>", "", text, flags=re.DOTALL)
    text = re.sub(r"<[^>]+>", "", text)

    # Remove tables {| ... |}
    text = re.sub(r"\{\|.*?\|\}", "", text, flags=re.DOTALL)

    # Convert wikilinks [[target|display]] -> display, [[target]] -> target
    text = re.sub(r"\[\[[^\]]*?\|([^\]]*?)\]\]", r"\1", text)
    text = re.sub(r"\[\[([^\]]*?)\]\]", r"\1", text)

    # Remove external links [url text] -> text
    text = re.sub(r"\[https?://\S+\s+([^\]]+)\]", r"\1", text)
    text = re.sub(r"\[https?://\S+\]", "", text)

    # Convert headers === Text === -> Text
    text = re.sub(r"^(=+)\s*(.+?)\s*\1\s*$", r"\2", text, flags=re.MULTILINE)

    # Remove bold/italic markers
    text = re.sub(r"'{2,5}", "", text)

    # Remove magic words / behavior switches
    text = re.sub(r"__[A-Z]+__", "", text)

    # Clean up whitespace
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)

    return text.strip()


def _fetch_parsed_html(title):
    """Fetch rendered HTML for a page via MediaWiki API and extract plaintext.

    Used as fallback for template-heavy pages where wikitext stripping yields nothing.
    Also follows redirects to other namespaces (e.g. Necyklopedie:Necyklokrám).
    """
    try:
        params = urllib.parse.urlencode({
            "action": "parse", "page": title, "prop": "text", "format": "json",
        })
        url = f"{API_URL}?{params}"
        req = urllib.request.Request(url, headers={"User-Agent": UA})
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        html = data.get("parse", {}).get("text", {}).get("*", "")
        if not html:
            return ""
        # Strip HTML tags, keeping text
        text = re.sub(r"<style[^>]*>.*?</style>", "", html, flags=re.DOTALL)
        text = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
        text = re.sub(r"<[^>]+>", " ", text)
        # Decode HTML entities
        import html as html_mod
        text = html_mod.unescape(text)
        # Clean whitespace
        text = re.sub(r"[ \t]+", " ", text)
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text.strip()
    except Exception:  # noqa: BLE001 - best-effort fallback (network, HTTP, JSON); "" skips the page
        return ""


# Minimum plaintext length from wikitext before trying HTML fallback
MIN_WIKITEXT_LEN = 20
# Minimum length for HTML fallback result to be accepted
MIN_HTML_LEN = 50
# Max pages to fetch via HTML fallback (rate-limited)
MAX_HTML_FALLBACKS = 200


def parse_dump():
    """Parse XML dump and extract articles as JSONL.

    Pages whose wikitext converts to very short plaintext (template-heavy pages)
    get a second chance via the MediaWiki parse API which renders templates.
    """
    print(f"Parsuju {DUMP_FILE} — převádím wikisyntaxi na čistý text...", file=sys.stderr)
    tree = ET.parse(DUMP_FILE)
    root = tree.getroot()

    total = 0
    skipped = 0
    html_fallback_candidates = []

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for page in root.findall(".//mw:page", MW_NS):
            # Skip redirects — but collect redirect targets for HTML fallback
            redirect_el = page.find("mw:redirect", MW_NS)
            if redirect_el is not None:
                title = page.findtext("mw:title", "", MW_NS)
                pageid = int(page.findtext("mw:id", "", MW_NS))
                target = redirect_el.get("title", "")
                # If redirect points to a non-main namespace, try HTML fallback
                if ":" in target and not target.startswith("Kategorie:"):
                    html_fallback_candidates.append((pageid, title))
                skipped += 1
                continue

            title = page.findtext("mw:title", "", MW_NS)
            pageid = int(page.findtext("mw:id", "", MW_NS))
            wikitext = page.findtext(".//mw:text", "", MW_NS)

            plaintext = wikitext_to_plaintext(wikitext)

            if len(plaintext) < MIN_WIKITEXT_LEN:
                # Template-heavy page — candidate for HTML fallback
                html_fallback_candidates.append((pageid, title))
                continue

            record = {"pageid": pageid, "title": title, "text": plaintext}
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
            total += 1

        print(f"\r[{total + skipped}] Uloženo {total}, přeskočeno {skipped}...", end="", file=sys.stderr)

    # Phase 2: HTML fallback for template-heavy pages
    if html_fallback_candidates:
        candidates = html_fallback_candidates[:MAX_HTML_FALLBACKS]
        print(f"\nZkouším HTML fallback pro {len(candidates)} šablonových stránek...", file=sys.stderr)
        fetched = 0
        with open(OUTPUT_FILE, "a", encoding="utf-8") as f:
            for i, (pageid, title) in enumerate(candidates):
                plaintext = _fetch_parsed_html(title)
                if len(plaintext) >= MIN_HTML_LEN:
                    record = {"pageid": pageid, "title": title, "text": plaintext}
                    f.write(json.dumps(record, ensure_ascii=False) + "\n")
                    fetched += 1
                    total += 1
                else:
                    skipped += 1
                if (i + 1) % 10 == 0:
                    print(f"\r  [{i + 1}/{len(candidates)}] Získáno {fetched}...", end="", file=sys.stderr)
                time.sleep(0.3)  # rate limit
        print(f"\n  HTML fallback: {fetched} stránek doplněno.", file=sys.stderr)

    print(f"Hotovo. {total} článků uloženo do {OUTPUT_FILE}. Pravda je připravena k indexaci.", file=sys.stderr)


def scrape():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Phase 1: enumerate titles
    titles = get_all_titles()

    # Phase 2: bulk export
    bulk_export(titles)

    # Phase 3: parse to plaintext JSONL
    parse_dump()


if __name__ == "__main__":
    scrape()