"""MediaWiki XML 'pages-articles' / 'pages-meta-current' / 'abstract' source.

Phase IV (2006+) Wikipedia dumps switched from MySQL extended INSERT
syntax to a custom XML schema. These sources stream that XML and yield
Documents the same way the WikipediaSqlDump source yields rows from
the 2003 cur table — same URI scheme, same edge extraction, same
shard / resume contract — so downstream ingest, sharding, and resume
logic Just Works.

Two source classes:

* `WikipediaXmlDump` — pages-articles.xml.bz2 (cur snapshot, main NS) or
  pages-meta-current.xml.bz2 (cur snapshot, all NS) or
  pages-meta-history.xml.bz2 (full history; pass multi_revision=True).
  One Document per page (or per revision in history mode).

* `WikipediaAbstractDump` — abstract.xml (uncompressed). Each `<doc>`
  has <title>, <url>, <abstract>: short pre-distilled summaries of every
  article. Useful as cheap retrieval targets for FTS5 queries.
"""

from __future__ import annotations

import bz2
import gzip
import re
from pathlib import Path
from typing import IO, Iterator
from xml.etree.ElementTree import Element, iterparse

from aborist.document import Document, Edge
from aborist.source import Source
from aborist.sources.wikipedia import _extract_wikilinks


# MediaWiki export schema declares an xmlns; ElementTree returns tag names
# as `{http://www.mediawiki.org/xml/export-0.10/}page`. Strip that prefix
# off so the rest of the parser doesn't have to care which schema version
# any given dump uses.
_NS_RE = re.compile(r"^\{[^}]+\}")

# Body that begins with `#REDIRECT [[Target]]` is a redirect — skip.
# Real dumps also use the explicit <redirect title="..."/> element, but
# older or imported pages sometimes lack that and only carry the body
# directive; we belt-and-suspenders both.
_REDIRECT_RE = re.compile(r"\s*#REDIRECT\b", re.IGNORECASE)

NAMESPACE_MAIN = 0


# Title-prefix-encoded namespaces. Only used as a fallback filter when the
# export schema doesn't declare <ns> per page. Excludes Wiktionary-style
# language prefixes (those don't appear in enwiki main namespace anyway).
_NON_MAIN_NS_PREFIXES = frozenset({
    "Talk", "User", "User talk", "Wikipedia", "Wikipedia talk",
    "File", "File talk", "MediaWiki", "MediaWiki talk",
    "Template", "Template talk", "Help", "Help talk",
    "Category", "Category talk", "Portal", "Portal talk",
    "Book", "Book talk", "Draft", "Draft talk",
    "Module", "Module talk", "Image", "Image talk",
})


def _strip_ns(tag: str) -> str:
    return _NS_RE.sub("", tag)


def _open_streaming(path: Path) -> IO[bytes]:
    """Open a file as a binary byte stream, transparently decompressing."""
    p = str(path)
    if p.endswith(".bz2"):
        return bz2.open(p, "rb")
    if p.endswith(".gz"):
        return gzip.open(p, "rb")
    return open(p, "rb")


class WikipediaXmlDump(Source):
    """Iterates a MediaWiki XML dump.

    Phase IV format: `<mediawiki>` root containing one `<siteinfo>` plus
    many `<page>` children. Each `<page>` has `<title>`, `<ns>`, `<id>`,
    optional `<redirect>`, and one or more `<revision>` children.

    For pages-articles or pages-meta-current dumps, each page has exactly
    one `<revision>` (the current one). For pages-meta-history dumps each
    page has many. By default we emit the LAST revision per page (matches
    cur semantics); pass `multi_revision=True` to emit every revision
    (matches old semantics, with `supersedes` chaining at the URI level
    via aborist's existing prior-document detection).
    """

    source_type = "wikipedia_xml"

    def __init__(
        self,
        path: str | Path,
        *,
        namespace: int = NAMESPACE_MAIN,
        base_uri: str = "https://en.wikipedia.org/wiki/",
        shard: tuple[int, int] | None = None,
        start_id: int = 0,
        multi_revision: bool = False,
    ):
        self.path = Path(path)
        self.namespace = namespace
        self.base_uri = base_uri
        if shard is not None:
            rank, total = shard
            if not (0 <= rank < total) or total < 1:
                raise ValueError(
                    f"invalid shard {shard}: need 0 <= rank < total, total >= 1"
                )
        self.shard = shard
        # Resume support — same contract as WikipediaSqlDump.
        self.start_id = start_id
        self.last_id: int = start_id
        self.multi_revision = multi_revision

    def iter_documents(self) -> Iterator[Document]:
        rank, total = (0, 1) if self.shard is None else self.shard
        idx = 0
        with _open_streaming(self.path) as f:
            context = iter(iterparse(f, events=("start", "end")))
            try:
                _, root = next(context)
            except StopIteration:
                return
            for event, elem in context:
                if event != "end":
                    continue
                tag = _strip_ns(elem.tag)
                if tag != "page":
                    # Free siteinfo/etc immediately — no useful data downstream.
                    if tag in ("siteinfo", "namespaces", "namespace"):
                        elem.clear()
                    continue
                # Stride filter applied at the page index. Every shard sees
                # the full XML stream but only emits its share — matches the
                # SQL source's behavior so ingest counters line up.
                if idx % total != rank:
                    idx += 1
                    elem.clear()
                    root.clear()
                    continue
                idx += 1
                yield from self._page_to_docs(elem)
                elem.clear()
                root.clear()

    def _page_to_docs(self, page: Element) -> Iterator[Document]:
        title: str | None = None
        ns_text: str | None = None
        page_id: int = 0
        is_redirect = False
        revisions: list[Element] = []
        for child in page:
            ctag = _strip_ns(child.tag)
            if ctag == "title":
                title = (child.text or "").strip() or None
            elif ctag == "ns":
                ns_text = (child.text or "").strip()
            elif ctag == "id":
                try:
                    page_id = int((child.text or "0").strip())
                except ValueError:
                    pass
            elif ctag == "redirect":
                # Self-closing element with title attr. Body may also start
                # with #REDIRECT — both flag the same condition.
                is_redirect = True
            elif ctag == "revision":
                revisions.append(child)

        try:
            ns_int = int(ns_text) if ns_text is not None else None
        except ValueError:
            ns_int = None
        # Older export schemas (pre-0.4) had no <ns> child element under
        # <page>; namespace was encoded into the title prefix only. The
        # `pages-articles.xml.bz2` dump is documented to contain just the
        # main namespace plus redirects, so trust the file when <ns> is
        # absent. For files that include sister namespaces (e.g.
        # `pages-meta-current.xml.bz2`), filter on the title prefix as
        # a safety belt — any title starting `Talk:` / `User:` /
        # `Wikipedia:` / `File:` / `Template:` / `Help:` / `Category:` /
        # `Portal:` / `Book:` / `Draft:` / `MediaWiki:` is non-main.
        if ns_int is None:
            if title and ":" in title:
                head = title.split(":", 1)[0]
                if head in _NON_MAIN_NS_PREFIXES:
                    return
        elif ns_int != self.namespace:
            return
        if is_redirect:
            return
        if not title or not revisions:
            return

        # Resume fast-forward.
        if self.start_id and page_id <= self.start_id:
            return
        if page_id > self.last_id:
            self.last_id = page_id

        revs = revisions if self.multi_revision else revisions[-1:]
        uri = self.base_uri + title.replace(" ", "_")
        for rev in revs:
            text = ""
            timestamp = ""
            rev_id = 0
            for child in rev:
                ctag = _strip_ns(child.tag)
                if ctag == "text":
                    text = child.text or ""
                elif ctag == "timestamp":
                    timestamp = (child.text or "").strip()
                elif ctag == "id":
                    try:
                        rev_id = int((child.text or "0").strip())
                    except ValueError:
                        pass
            if not text:
                continue
            if _REDIRECT_RE.match(text):
                continue
            edges = _extract_wikilinks(text, self.base_uri)
            extra: dict = {"page_id": str(page_id)}
            if timestamp:
                extra["timestamp"] = timestamp
            if rev_id:
                extra["rev_id"] = str(rev_id)
            yield Document(
                uri=uri,
                content=text,
                source_type=self.source_type,
                title=title,
                edges=edges,
                extra=extra,
            )


class WikipediaAbstractDump(Source):
    """Iterates a MediaWiki `abstract.xml` dump.

    Format:
        <feed>
          <doc>
            <title>Wikipedia: Anarchism</title>
            <url>http://en.wikipedia.org/wiki/Anarchism</url>
            <abstract>Anarchism is a political philosophy...</abstract>
            <links>...</links>
          </doc>
          ...
        </feed>

    Each <doc> already contains a one-paragraph summary — pre-distilled
    by the dump generator. Ingesting these alongside (or instead of) full
    pages-articles gives FTS5 a high-signal corpus at ~1/100th the bytes.
    """

    source_type = "wikipedia_abstract"

    def __init__(self, path: str | Path):
        self.path = Path(path)
        # Resume contract: abstract docs have no stable numeric id, so
        # start_id / last_id are no-ops here. Kept for interface symmetry.
        self.start_id = 0
        self.last_id = 0

    def iter_documents(self) -> Iterator[Document]:
        with _open_streaming(self.path) as f:
            context = iter(iterparse(f, events=("start", "end")))
            try:
                _, root = next(context)
            except StopIteration:
                return
            for event, elem in context:
                if event != "end":
                    continue
                tag = _strip_ns(elem.tag)
                if tag != "doc":
                    continue
                title = ""
                url = ""
                abstract = ""
                links: list[str] = []
                for child in elem:
                    ctag = _strip_ns(child.tag)
                    if ctag == "title":
                        title = (child.text or "").strip()
                    elif ctag == "url":
                        url = (child.text or "").strip()
                    elif ctag == "abstract":
                        abstract = (child.text or "").strip()
                    elif ctag == "links":
                        for sublink in child:
                            for sub in sublink:
                                if _strip_ns(sub.tag) == "link":
                                    href = (sub.text or "").strip()
                                    if href:
                                        links.append(href)
                elem.clear()
                root.clear()
                if not url or not abstract:
                    continue
                # Abstract feed prepends "Wikipedia: " to titles. Strip it.
                clean_title = title
                if clean_title.startswith("Wikipedia: "):
                    clean_title = clean_title[len("Wikipedia: "):]
                edges = [
                    Edge(edge_type="abstract_link", dst_uri=u) for u in links
                ]
                yield Document(
                    uri=url,
                    content=abstract,
                    source_type=self.source_type,
                    title=clean_title or None,
                    edges=edges,
                )
