"""
common/arxiv.py — first-class ARXIV external-source type.

Operator directive (2026-06-24): "map it down and be preparing for this specific
type of external sources everywhere in the parsers". When ingestion / web-fetch /
the safekeeping-mirror encounters an arxiv link ANYWHERE (a curated list, a blog
post, a bare id, a DOI), it must be handled as a TYPED ARXIV source — a structured
"card" with all the arxiv attributes — and its content extracted from the e-print
(LaTeX source) or the rendered PDF (whichever lets us best recover formulae /
graphs / tables).

This module is the FOUNDATION:

  1. recognise()      — robust arxiv-id / URL / DOI matcher (new- AND old-style ids,
                        abs|pdf|format paths, export.arxiv.org, 10.48550/arXiv.<id>,
                        bare 2501.12948 / cond-mat/0703012).
  2. ArxivCard        — the structured card built from arxiv's OFFICIAL metadata API
                        (export.arxiv.org/api/query, the Atom feed — NEVER HTML scraping).
  3. fetch_card()     — id → metadata card (respects arxiv API rate limits / ToS).
  4. eprint / pdf urls + fetch_eprint() / fetch_pdf() — pull the LaTeX source tarball
                        and/or the rendered PDF into a bytes payload.
  5. ingest_arxiv()   — the typed-source submission: card + body → the v2 pipeline as
                        source_type="arxiv", provider="arxiv".

DEEP EXTRACTION — TeX path is BUILT (Phase 1, #157):
  • extract_eprint() parses the author's LaTeX e-print with **pylatexenc** (a real
    parser, not regex) into STRUCTURED, TYPED content — body + formulae[] + tables[]
    + figures[] + references[] — that is CARRIED THROUGH the pipeline, never flattened
    into prose. ingest_arxiv() prefers this structured body and rides the typed arrays
    under metadata.arxiv.extraction for the chunker/extractor (Phase 3).
  • extract_pdf() (fallback when the e-print is absent/unparseable) is still a Phase-2
    seam — see EXTRACTION-MAP + andrea/infrastructure#157.

Deps: stdlib urllib + a tolerant XML parse of the Atom feed; **pylatexenc** for deep
TeX extraction (graceful pragmatic-regex fallback when pylatexenc is absent, so the
recognition / fetch / card path stays dependency-light).
"""
from __future__ import annotations

import logging
import os
import re
import tarfile
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass, field, asdict
from datetime import datetime, timezone
from io import BytesIO
from typing import Any, Optional
from xml.etree import ElementTree as ET

logger = logging.getLogger(__name__)

# ── arxiv endpoints (source-of-truth; override via env, never hardcode elsewhere) ──
ARXIV_API_URL   = os.getenv("ARXIV_API_URL",   "https://export.arxiv.org/api/query")
ARXIV_ABS_BASE  = os.getenv("ARXIV_ABS_BASE",  "https://arxiv.org/abs")
ARXIV_PDF_BASE  = os.getenv("ARXIV_PDF_BASE",  "https://arxiv.org/pdf")
# The e-print (LaTeX source tarball) endpoint. arxiv serves it at /e-print/<id>.
ARXIV_EPRINT_BASE = os.getenv("ARXIV_EPRINT_BASE", "https://arxiv.org/e-print")
ARXIV_DOI_PREFIX  = "10.48550/arxiv."          # lower-cased for matching

# arxiv ToS: be gentle. One request per ~3 s, identify ourselves.
ARXIV_MIN_INTERVAL = float(os.getenv("ARXIV_MIN_INTERVAL", "3.0"))
ARXIV_USER_AGENT   = os.getenv(
    "ARXIV_USER_AGENT",
    "AndreaIngest/1.0 (+https://projectandrea.com; mailto:avesnin@protonmail.ch)")
ARXIV_HTTP_TIMEOUT = int(os.getenv("ARXIV_HTTP_TIMEOUT", "30"))

SOURCE_TYPE = "arxiv"
PROVIDER    = "arxiv"

# Atom / arxiv / OpenSearch XML namespaces (the API returns an Atom feed).
_NS = {
    "atom":    "http://www.w3.org/2005/Atom",
    "arxiv":   "http://arxiv.org/schemas/atom",
    "opensearch": "http://a9.com/-/spec/opensearch/1.1/",
}


# ════════════════════════════════════════════════════════════════════════════════
#  1. RECOGNITION — the robust matcher
# ════════════════════════════════════════════════════════════════════════════════
#
# arxiv ids come in two eras:
#   • NEW (since 2007-04):  YYMM.NNNNN   e.g. 2501.12948, 0704.0001  (4- or 5-digit
#     sequence; 5 digits since 2015-01). Optionally a version suffix vN.
#   • OLD (pre-2007-04):    archive[.subclass]/YYMMNNN   e.g. cond-mat/0703012,
#     math.AG/0309136, hep-th/9901001. Optionally vN.
#
# A version suffix (vN) is OPTIONAL everywhere and, when absent, means "latest".
#
# We accept the id bare, or wrapped in any of the canonical URL shapes, or as the
# arxiv DOI (10.48550/arXiv.<id>). HTML scraping is never used — recognition only
# extracts the id; the CARD always comes from the official API.

# New-style: 4-digit YYMM + dot + 4-or-5-digit sequence, optional vN.
_NEW_ID = r"\d{4}\.\d{4,5}(?:v\d+)?"
# Old-style archive token: letters, optional dot+subclass (letters/dash), then
# /YYMMNNN, optional vN. Archive set is open-ended but always lower-case letters.
_OLD_ID = r"[a-z][a-z\-]+(?:\.[A-Z]{2})?/\d{7}(?:v\d+)?"

# Bare-id matcher (for "2501.12948" / "cond-mat/0703012" appearing in free text).
_RE_BARE = re.compile(rf"(?<![\w./-])(?:arXiv:)?({_NEW_ID}|{_OLD_ID})(?![\w])",
                      re.IGNORECASE)

# URL/DOI host+path shapes we route on. The id is captured group 1.
_URL_PATTERNS = [
    # arxiv.org/abs/<id> , /pdf/<id> , /pdf/<id>.pdf , /format/<id> , /e-print/<id>
    re.compile(rf"arxiv\.org/(?:abs|pdf|format|e-print|html)/({_NEW_ID}|{_OLD_ID})(?:\.pdf)?",
               re.IGNORECASE),
    # export.arxiv.org variants (the API host; abs/pdf also resolve there)
    re.compile(rf"export\.arxiv\.org/(?:abs|pdf|e-print)/({_NEW_ID}|{_OLD_ID})(?:\.pdf)?",
               re.IGNORECASE),
    # DOI form: 10.48550/arXiv.<id>  (doi.org/… or bare doi:)
    re.compile(rf"10\.48550/arxiv\.({_NEW_ID}|{_OLD_ID})", re.IGNORECASE),
]


def normalize_id(raw: str) -> tuple[str, Optional[str]]:
    """Split a raw arxiv id into (base_id, version).

    >>> normalize_id("2501.12948v3")  -> ("2501.12948", "v3")
    >>> normalize_id("cond-mat/0703012") -> ("cond-mat/0703012", None)

    The base id is lower-cased ONLY on its archive token for old-style ids (new-style
    ids have no letters). Version is returned as the literal "vN" or None.
    """
    raw = raw.strip()
    m = re.search(r"(v\d+)$", raw)
    version = None
    if m:
        version = m.group(1)
        raw = raw[: -len(version)]
    # old-style archive token is conventionally lower-case; new-style is digits only.
    if "/" in raw:
        arch, num = raw.split("/", 1)
        raw = f"{arch.lower()}/{num}"
    return raw, version


def recognize(text: str) -> Optional[str]:
    """Return the FIRST arxiv id found anywhere in `text` (a URL, a DOI, or a bare
    id), or None. Returns the id WITH its version suffix if one was present, so the
    caller can fetch that exact version. Use this as the routing predicate in the
    dissector / web-fetch: a non-None result means "route to the arxiv handler".

    Recognition order: URL/DOI shapes first (most specific), then bare-id scan.
    """
    if not text:
        return None
    for pat in _URL_PATTERNS:
        m = pat.search(text)
        if m:
            base, ver = normalize_id(m.group(1))
            return base + (ver or "")
    m = _RE_BARE.search(text)
    if m:
        base, ver = normalize_id(m.group(1))
        return base + (ver or "")
    return None


def recognize_all(text: str) -> list[str]:
    """Return EVERY distinct arxiv id referenced in `text`, in first-seen order.

    For curated lists / blog posts that link many papers (the motivating example:
    a magazine post linking a dozen arxiv abstracts). Deduplicated on base id +
    explicit version.
    """
    seen: list[str] = []
    if not text:
        return seen

    def _add(raw: str) -> None:
        base, ver = normalize_id(raw)
        full = base + (ver or "")
        if full not in seen:
            seen.append(full)

    for pat in _URL_PATTERNS:
        for m in pat.finditer(text):
            _add(m.group(1))
    for m in _RE_BARE.finditer(text):
        _add(m.group(1))
    return seen


def is_arxiv(text: str) -> bool:
    """Cheap boolean predicate for routing — True iff `text` references an arxiv id."""
    return recognize(text) is not None


# ════════════════════════════════════════════════════════════════════════════════
#  2. THE ARXIV CARD — structured attributes from the OFFICIAL API
# ════════════════════════════════════════════════════════════════════════════════

@dataclass
class ArxivAuthor:
    name: str
    affiliation: Optional[str] = None      # present only when the submitter tagged it


@dataclass
class ArxivCard:
    """The structured card for one arxiv paper, built from the API Atom entry.

    Stored verbatim into source_documents.metadata.arxiv (under the typed source) so
    every downstream consumer — search filters, the citation spine, the deep
    extractor — reads the same canonical attributes. `arxiv_id`+`version` form the
    natural key; `provider` is always "arxiv".
    """
    arxiv_id: str                                  # base id, no version (e.g. 2501.12948)
    version: Optional[str] = None                  # "v3" or None (= latest at fetch time)
    title: str = ""
    authors: list[ArxivAuthor] = field(default_factory=list)
    abstract: str = ""
    primary_category: str = ""                     # e.g. cs.LG
    categories: list[str] = field(default_factory=list)   # primary + cross-list
    comments: Optional[str] = None                 # e.g. "12 pages, 5 figures"
    doi: Optional[str] = None                       # author-supplied DOI (NOT the 10.48550 one)
    journal_ref: Optional[str] = None
    published: Optional[str] = None                # submitted date (ISO8601)
    updated: Optional[str] = None                  # last-updated date (ISO8601)
    license: Optional[str] = None                  # license URL when present
    abs_url: str = ""
    pdf_url: str = ""
    eprint_url: str = ""
    provider: str = PROVIDER
    fetched_at: str = ""

    def to_metadata(self, title_override: Optional[str] = None) -> dict[str, Any]:
        """Render the pipeline `metadata` dict for this card. The pipeline keys
        (`title`, `source_type`, `provider`) sit at top level so storage/worker.py
        lands them on the source_documents row; the full card is nested under
        `arxiv` for lossless downstream use."""
        return {
            "title":       title_override or self.title,
            "source_type": SOURCE_TYPE,
            "provider":    PROVIDER,
            "arxiv":       asdict(self),
        }


def _txt(el: Optional[ET.Element]) -> str:
    return (el.text or "").strip() if el is not None else ""


def _parse_entry(entry: ET.Element) -> ArxivCard:
    """Parse one Atom <entry> into an ArxivCard. Tolerant of missing optional fields."""
    # id is a URL like http://arxiv.org/abs/2501.12948v3 — derive id+version from it.
    raw_id = _txt(entry.find("atom:id", _NS))
    id_with_ver = recognize(raw_id) or raw_id.rsplit("/", 1)[-1]
    base_id, version = normalize_id(id_with_ver)

    authors: list[ArxivAuthor] = []
    for a in entry.findall("atom:author", _NS):
        name = _txt(a.find("atom:name", _NS))
        aff = a.find("arxiv:affiliation", _NS)
        authors.append(ArxivAuthor(name=name,
                                   affiliation=_txt(aff) or None if aff is not None else None))

    primary = entry.find("arxiv:primary_category", _NS)
    primary_cat = primary.get("term", "") if primary is not None else ""
    cats: list[str] = []
    if primary_cat:
        cats.append(primary_cat)
    for c in entry.findall("atom:category", _NS):
        term = c.get("term", "")
        if term and term not in cats:
            cats.append(term)

    # license / pdf links live in <link> elements with rel/type/title attrs.
    pdf_url = ""
    license_url = None
    for ln in entry.findall("atom:link", _NS):
        if ln.get("title") == "pdf" or ln.get("type") == "application/pdf":
            pdf_url = ln.get("href", "") or pdf_url
        if ln.get("rel") == "license":
            license_url = ln.get("href") or license_url

    return ArxivCard(
        arxiv_id=base_id,
        version=version,
        title=re.sub(r"\s+", " ", _txt(entry.find("atom:title", _NS))).strip(),
        authors=authors,
        abstract=re.sub(r"\s+", " ", _txt(entry.find("atom:summary", _NS))).strip(),
        primary_category=primary_cat,
        categories=cats,
        comments=_txt(entry.find("arxiv:comment", _NS)) or None,
        doi=_txt(entry.find("arxiv:doi", _NS)) or None,
        journal_ref=_txt(entry.find("arxiv:journal_ref", _NS)) or None,
        published=_txt(entry.find("atom:published", _NS)) or None,
        updated=_txt(entry.find("atom:updated", _NS)) or None,
        license=license_url,
        abs_url=abs_url(id_with_ver),
        pdf_url=pdf_url or pdf_url_for(id_with_ver),
        eprint_url=eprint_url(id_with_ver),
        provider=PROVIDER,
        fetched_at=datetime.now(timezone.utc).isoformat(),
    )


# ════════════════════════════════════════════════════════════════════════════════
#  3 & 4. FETCH — metadata card + e-print / pdf (rate-limited, ToS-respecting)
# ════════════════════════════════════════════════════════════════════════════════

_last_request_ts = 0.0


def _throttle() -> None:
    """Enforce the arxiv minimum inter-request interval (ToS / be-a-good-citizen)."""
    global _last_request_ts
    wait = ARXIV_MIN_INTERVAL - (time.monotonic() - _last_request_ts)
    if wait > 0:
        time.sleep(wait)
    _last_request_ts = time.monotonic()


def _http_get(url: str, accept: str = "*/*") -> bytes:
    req = urllib.request.Request(
        url, headers={"User-Agent": ARXIV_USER_AGENT, "Accept": accept})
    _throttle()
    with urllib.request.urlopen(req, timeout=ARXIV_HTTP_TIMEOUT) as resp:
        return resp.read()


def abs_url(id_with_ver: str) -> str:
    return f"{ARXIV_ABS_BASE}/{id_with_ver}"


def pdf_url_for(id_with_ver: str) -> str:
    return f"{ARXIV_PDF_BASE}/{id_with_ver}"


def eprint_url(id_with_ver: str) -> str:
    return f"{ARXIV_EPRINT_BASE}/{id_with_ver}"


def fetch_card(arxiv_id: str) -> ArxivCard:
    """Fetch the metadata CARD for one arxiv id from the OFFICIAL API (Atom feed).

    `arxiv_id` may carry a version (e.g. 2501.12948v3); the API `id_list` honours it.
    Raises LookupError if the id is unknown / the feed has no entry.
    """
    base, version = normalize_id(arxiv_id)
    id_for_api = base + (version or "")
    qs = urllib.parse.urlencode({"id_list": id_for_api, "max_results": 1})
    raw = _http_get(f"{ARXIV_API_URL}?{qs}", accept="application/atom+xml")
    root = ET.fromstring(raw)
    entries = root.findall("atom:entry", _NS)
    # arxiv returns a sentinel "Error" entry (title="Error") for an unknown id.
    for e in entries:
        title = _txt(e.find("atom:title", _NS))
        eid = _txt(e.find("atom:id", _NS))
        if title == "Error" or "api/errors" in eid:
            raise LookupError(f"arxiv id not found: {arxiv_id!r} ({_txt(e.find('atom:summary', _NS))})")
    if not entries:
        raise LookupError(f"arxiv id not found (empty feed): {arxiv_id!r}")
    return _parse_entry(entries[0])


def search_cards(query: str, max_results: int = 10) -> list[ArxivCard]:
    """Free-text arxiv API search (`search_query`) → cards. Not on the ingest hot
    path; provided for tooling / discovery (e.g. resolving a title to an id)."""
    qs = urllib.parse.urlencode(
        {"search_query": query, "max_results": max_results,
         "sortBy": "relevance", "sortOrder": "descending"})
    raw = _http_get(f"{ARXIV_API_URL}?{qs}", accept="application/atom+xml")
    root = ET.fromstring(raw)
    return [_parse_entry(e) for e in root.findall("atom:entry", _NS)]


def fetch_eprint(arxiv_id: str) -> Optional[bytes]:
    """Fetch the e-print (LaTeX source) tarball bytes for an id, or None if absent
    (withdrawn paper, or a paper with no source — only a publisher PDF). The bytes
    are the raw tar(.gz) blob; `iter_eprint_files` unpacks it.

    NB: a few e-prints are a single .pdf or a gzipped single file, not a tar — the
    caller / extractor must sniff (see iter_eprint_files)."""
    try:
        return _http_get(eprint_url(arxiv_id), accept="application/x-eprint-tar")
    except urllib.error.HTTPError as e:           # 404 = no source on file
        if e.code == 404:
            logger.info("arxiv e-print absent for %s (HTTP 404) — PDF fallback", arxiv_id)
            return None
        raise


def fetch_pdf(arxiv_id: str) -> bytes:
    """Fetch the rendered PDF bytes for an id. Always available for a non-withdrawn
    paper; the FALLBACK extraction source when the e-print is absent/unparseable."""
    return _http_get(pdf_url_for(arxiv_id) + ".pdf", accept="application/pdf")


def iter_eprint_files(blob: bytes):
    """Yield (name, bytes) for each member of an e-print blob. Handles the common
    cases: a gzipped tar (most papers), a bare tar, or a single gzipped TeX file.
    A sniff-and-unpack convenience for the extractor seam — keeps the tar handling
    in ONE place. Best-effort: unrecognised blobs yield nothing.

    SEAM: the deep extractor consumes this to find the main .tex, \\input children,
    .bbl/.bib, and figure source files (.pdf/.eps/.png/.jpg) — see EXTRACTION-MAP."""
    if not blob:
        return
    bio = BytesIO(blob)
    try:
        with tarfile.open(fileobj=bio, mode="r:*") as tf:
            for m in tf.getmembers():
                if not m.isfile():
                    continue
                f = tf.extractfile(m)
                if f is not None:
                    yield m.name, f.read()
        return
    except tarfile.ReadError:
        pass
    # Not a tar — try a single gzip member.
    import gzip
    try:
        bio.seek(0)
        data = gzip.decompress(bio.read())
        yield "main.tex", data           # arxiv single-file source is conventionally TeX
    except OSError:
        logger.warning("e-print blob is neither tar nor gzip (%d bytes) — skipping", len(blob))


# ════════════════════════════════════════════════════════════════════════════════
#  EXTRACTION-MAP  (deep extraction — SEAMS, mapped not built)
# ════════════════════════════════════════════════════════════════════════════════
#
# WHAT we extract FROM WHICH source (PDF-vs-TeX strategy). Prefer TeX; PDF is the
# fallback. The functions below are the SEAMS where a later worker fills the parser.
#
#   ┌──────────────┬──────────────────────────┬───────────────────────────────────┐
#   │ Atom         │ PRIMARY (e-print / TeX)  │ FALLBACK (rendered PDF)            │
#   ├──────────────┼──────────────────────────┼───────────────────────────────────┤
#   │ formulae     │ LaTeX math verbatim      │ OCR / pdf-math (MathML, lossy)    │
#   │              │ ($…$, \\[ \\], equation, │ — only when no e-print            │
#   │              │ align) → keep LaTeX +    │                                   │
#   │              │ render MathML            │                                   │
#   │ tables       │ tabular / longtable /    │ table detection on the PDF        │
#   │              │ booktabs → structured    │ (camelot/tabula-class) — lossy    │
#   │              │ rows                     │                                   │
#   │ figures      │ figure source files in   │ rasterised page regions from PDF  │
#   │              │ the e-print (.pdf/.eps/  │                                   │
#   │              │ .png/.jpg) + \\caption   │                                   │
#   │ body text    │ TeX body, macros expanded│ PDF text layer (pdftotext-class)  │
#   │ references   │ .bbl / .bib / \\bibitem  │ PDF reference section parse       │
#   └──────────────┴──────────────────────────┴───────────────────────────────────┘
#
# Output contract — extract_eprint() (Phase 1, BUILT) returns EXACTLY this so the
# chunker/extractor carry the TYPED content through (each typed atom keeps verbatim
# LaTeX; nothing is flattened into `body`):
#   {
#     "body":      str,                 # readable prose, section structure preserved,
#                                       #   commands stripped; typed envs NOT duplicated
#     "formulae":  [{"latex": str,      #   VERBATIM source slice
#                    "mathml": None,    #   Phase 2 (latexml/converter)
#                    "label": str|None,
#                    "section_context": str|None,
#                    "env": str}, ...], #   inline|display|equation|align|…
#     "tables":    [{"caption": str|None, "label": str|None,
#                    "rows": [[cell, ...], ...],   # parsed on &/\\
#                    "latex": str,      #   VERBATIM source slice
#                    "section_context": str|None}, ...],
#     "figures":   [{"caption": str|None, "label": str|None,
#                    "graphics": [str, ...],       # \includegraphics targets
#                    "latex": str,      #   VERBATIM source slice
#                    "section_context": str|None}, ...],
#     "references":[{"key": str, "authors": str|None, "title": str|None,
#                    "year": str|None, "venue": str|None,
#                    "arxiv_id": str|None, "doi": str|None}, ...],
#     "source":    "tex" | "pdf",       # which path produced this
#     "main_tex":  str|None,            # the main .tex key (\documentclass+document)
#     "warnings":  [str, ...],          # unresolved \input, parser fallbacks, …
#   }
# The extractor/storage seams may add "image_key" onto figures and "mathml" onto
# formulae when those are produced (Phase 2/3) — additive, not replacing.
#
# TOUCHPOINTS for the deep worker (concrete file+worker seams):
#   • common/arxiv.py        (HERE)   extract_eprint() BUILT (Phase 1) / extract_pdf()
#                                     Phase-2 seam.
#   • dissector/chunker.py            consume body+formulae+tables so a formula/table
#                                     is its own atom (don't split mid-equation); carry
#                                     the latex/mathml onto the chunk so retrieval can
#                                     surface the math. (chunker is the chunk-boundary
#                                     authority — see dissector/pipeline.py semantic_chunk.)
#   • extractor/worker.py             figures→image atoms in Garage (image_key), keyed
#                                     under pipeline/<ctx>/<doc>/figures/; references→
#                                     citation edges (each ref arxiv_id/doi is a link
#                                     the citation spine dereferences).
#   • storage/worker.py               persist figure image_keys + formula/table arrays
#                                     onto the chunk rows (already carries chunk dicts
#                                     verbatim — just add the new keys to write_chunk).


# ════════════════════════════════════════════════════════════════════════════════
#  DEEP TeX EXTRACTION — Phase 1 (BUILT). The PRIMARY path: parse the author's LaTeX
#  source into STRUCTURED, TYPED content (formulae / tables / figures / references)
#  that is CARRIED THROUGH the pipeline, never flattened into prose.
# ════════════════════════════════════════════════════════════════════════════════
#
# Implemented with **pylatexenc** (pure-python, LatexWalker) — a real LaTeX tokenizer/
# parser, NOT brittle regex. It understands nested environments, group nesting,
# comments, verbatim, and macro arguments, so `$…$` inside a `\caption{}` inside a
# `figure` is parsed correctly and a `\\` inside an inline math node never splits a
# tabular row. When pylatexenc is unavailable we degrade to a pragmatic env-aware
# scanner (see _PylatexencUnavailable handling) rather than failing the ingest.
#
# Math environments we treat as a single typed formula atom (display + numbered):
_MATH_ENVS = frozenset({
    "equation", "equation*", "align", "align*", "alignat", "alignat*",
    "gather", "gather*", "multline", "multline*", "eqnarray", "eqnarray*",
    "flalign", "flalign*", "math", "displaymath", "split", "cases", "array",
})
# Float environments that carry tables / figures.
_TABLE_ENVS = frozenset({"table", "table*", "tabular", "tabular*", "longtable",
                         "tabularx", "supertabular", "array"})
_FIGURE_ENVS = frozenset({"figure", "figure*", "wrapfigure", "subfigure"})
# Environments whose body is NOT prose and must be skipped by the prose walker
# (their typed content is harvested separately).
_NON_PROSE_ENVS = _MATH_ENVS | _TABLE_ENVS | _FIGURE_ENVS | frozenset(
    {"thebibliography", "abstract", "tabular", "tabular*"})

# pylatexenc availability is probed lazily and cached, so module import stays light.
_PYLATEXENC: Optional[bool] = None


def _pylatexenc_context():
    """Build (and memoise) a LatexContextDb that knows the arxiv-relevant macros take
    arguments: \\caption{}, \\includegraphics[]{} , \\label{}, \\input{}/\\include{},
    \\cite{}, \\bibitem[]{}, \\title{}, \\author{}, \\section*{} family. Without these
    the default context treats e.g. \\caption as zero-arg and the caption text leaks
    into the body. Returns (context_db, macrospec_module) or raises ImportError."""
    from pylatexenc.latexwalker import get_default_latex_context_db
    from pylatexenc.macrospec import std_macro

    db = get_default_latex_context_db()
    db.add_context_category(
        "andrea-arxiv",
        macros=[
            std_macro("caption", True, 1),       # \caption[short]{long}
            std_macro("includegraphics", True, 1),
            std_macro("label", False, 1),
            std_macro("ref", False, 1),
            std_macro("input", False, 1),
            std_macro("include", False, 1),
            std_macro("subfile", False, 1),
            std_macro("import", False, 2),
            std_macro("cite", True, 1),
            std_macro("citep", True, 1),
            std_macro("citet", True, 1),
            std_macro("bibitem", True, 1),
            std_macro("bibliography", False, 1),
            std_macro("title", False, 1),
            std_macro("author", False, 1),
        ],
        prepend=True,
    )
    return db


def _verbatim(src: str, node) -> str:
    """The EXACT source slice for a node — the load-bearing primitive that keeps math
    / tables / figures VERBATIM (we never re-serialise LaTeX, we slice the original)."""
    return src[node.pos: node.pos + node.len]


def _env_body_verbatim(src: str, env) -> str:
    """The verbatim source BETWEEN \\begin{env}…\\end{env} (the body only, excluding
    the begin/end and the env's own bracket/brace args)."""
    if not getattr(env, "nodelist", None):
        return ""
    first = env.nodelist[0]
    last = env.nodelist[-1]
    return src[first.pos: last.pos + last.len]


def _macro_arg_text(src: str, macro, idx: int = -1) -> Optional[str]:
    """Return the (brace-stripped) verbatim text of macro argument `idx` (default: the
    LAST mandatory arg — e.g. the file in \\includegraphics[opts]{file}, the key in
    \\label{key}). None if the macro has no such argument."""
    argd = getattr(macro, "nodeargd", None)
    if not argd or not argd.argnlist:
        return None
    args = [a for a in argd.argnlist if a is not None]
    if not args:
        return None
    arg = args[idx]
    raw = src[arg.pos: arg.pos + arg.len].strip()
    if len(raw) >= 2 and raw[0] in "{[" and raw[-1] in "}]":
        raw = raw[1:-1]
    return raw.strip()


def _to_text(latex_fragment: str) -> str:
    """Render a LaTeX fragment to readable plain prose (commands stripped to text).
    Used for captions and the prose body. Falls back to a light strip if latex2text
    is unavailable or chokes on a fragment."""
    try:
        from pylatexenc.latex2text import LatexNodes2Text
        return re.sub(r"[ \t]+", " ",
                      LatexNodes2Text().latex_to_text(latex_fragment)).strip()
    except Exception:                                    # noqa: BLE001
        # crude fallback: drop the most common commands, collapse whitespace
        s = re.sub(r"\\[a-zA-Z@]+\*?(\[[^\]]*\])?(\{[^{}]*\})?", " ", latex_fragment)
        s = s.replace("{", "").replace("}", "")
        return re.sub(r"\s+", " ", s).strip()


def _strip_tex_comments(text: str) -> str:
    """Remove TeX line-comments (`%` to EOL) while honouring an escaped `\\%`. Applied
    only on the pragmatic-fallback path; the pylatexenc path handles comments natively."""
    out = []
    for line in text.splitlines():
        i, n, res = 0, len(line), []
        while i < n:
            c = line[i]
            if c == "\\" and i + 1 < n:
                res.append(line[i:i + 2]); i += 2; continue
            if c == "%":
                break
            res.append(c); i += 1
        out.append("".join(res))
    return "\n".join(out)


# ── e-print file-set assembly + main-.tex discovery + \input resolution ─────────

def _eprint_file_map(blob_or_dir) -> dict[str, bytes]:
    """Normalise the e-print into {posix-relative-path: bytes}. Accepts:
      • bytes        — a raw e-print tar(.gz)/gzip blob (uses iter_eprint_files);
      • a dir path   — an already-unpacked e-print directory (walk it);
      • a dict       — already a {name: bytes} map (pass-through, e.g. tests)."""
    if isinstance(blob_or_dir, dict):
        return {k.lstrip("./"): v for k, v in blob_or_dir.items()}
    if isinstance(blob_or_dir, (bytes, bytearray)):
        files: dict[str, bytes] = {}
        for name, data in iter_eprint_files(bytes(blob_or_dir)):
            files[name.lstrip("./")] = data
        return files
    # a directory path
    root = str(blob_or_dir)
    files = {}
    for dirpath, _dirs, fnames in os.walk(root):
        for fn in fnames:
            full = os.path.join(dirpath, fn)
            rel = os.path.relpath(full, root).replace(os.sep, "/")
            try:
                with open(full, "rb") as fh:
                    files[rel] = fh.read()
            except OSError:
                continue
    return files


def _decode(data: bytes) -> str:
    """Best-effort decode of a TeX file — UTF-8, then latin-1 (old arxiv sources)."""
    for enc in ("utf-8", "latin-1"):
        try:
            return data.decode(enc)
        except UnicodeDecodeError:
            continue
    return data.decode("utf-8", errors="replace")


def _find_main_tex(files: dict[str, bytes]) -> Optional[str]:
    """Pick the main .tex: the one carrying BOTH \\documentclass and \\begin{document}.
    Tie-break toward a top-level (shallowest) file, then a conventional name
    (main/ms/paper/root/article), then the largest. Returns the file key or None."""
    candidates = []
    for name, data in files.items():
        if not name.lower().endswith(".tex"):
            continue
        text = _decode(data)
        score = 0
        if re.search(r"\\documentclass", text):
            score += 2
        if re.search(r"\\begin\s*\{document\}", text):
            score += 3
        if score == 0:
            continue
        depth = name.count("/")
        base = os.path.splitext(os.path.basename(name))[0].lower()
        name_bonus = 1 if base in {"main", "ms", "paper", "root", "article",
                                   "manuscript"} else 0
        candidates.append((score, name_bonus, -depth, len(data), name))
    if not candidates:
        # no \begin{document} anywhere (rare: a fragment-only source) — take the
        # largest .tex as a last resort so we still extract SOMETHING.
        texs = [(len(d), n) for n, d in files.items() if n.lower().endswith(".tex")]
        return max(texs)[1] if texs else None
    candidates.sort(reverse=True)
    return candidates[0][4]


def _resolve_input(arg: str, base_dir: str, files: dict[str, bytes]) -> Optional[str]:
    """Resolve an \\input/\\include target (extension often omitted; path relative to
    the including file's dir, but arxiv flattens — so also try the e-print root and a
    basename match). Returns the resolved file key or None."""
    arg = arg.strip().strip("\"")
    cands = []
    for stem in ({arg, arg + ".tex"} if not arg.lower().endswith(".tex")
                 else {arg}):
        # relative to including file
        if base_dir:
            cands.append(f"{base_dir}/{stem}".replace("//", "/"))
        cands.append(stem)                         # relative to root
    # normalise ./ and ../ lightly
    norm = []
    for c in cands:
        parts, stack = c.split("/"), []
        for p in parts:
            if p in ("", "."):
                continue
            if p == ".." and stack:
                stack.pop()
            else:
                stack.append(p)
        norm.append("/".join(stack))
    for c in norm:
        if c in files:
            return c
    # basename fallback (arxiv sometimes flattens subdirs)
    want = os.path.basename(norm[0]) if norm else None
    if want:
        for k in files:
            if os.path.basename(k) == want:
                return k
    return None


def _flatten_tex(main_key: str, files: dict[str, bytes],
                 warnings: list[str]) -> str:
    """Expand \\input/\\include/\\subfile recursively into one combined source string.
    Guards against cycles and missing children (recorded as warnings, never raised)."""
    seen: set[str] = set()

    def expand(key: str, depth: int) -> str:
        if key in seen:
            warnings.append(f"input cycle/repeat skipped: {key}")
            return ""
        if depth > 25:
            warnings.append(f"input nesting too deep at {key}")
            return ""
        seen.add(key)
        text = _decode(files.get(key, b""))
        base_dir = os.path.dirname(key)

        def _sub(m: "re.Match") -> str:
            child_arg = m.group(2)
            child = _resolve_input(child_arg, base_dir, files)
            if child is None:
                warnings.append(f"unresolved \\{m.group(1)}{{{child_arg}}} in {key}")
                return ""                              # drop the unresolved include
            return "\n" + expand(child, depth + 1) + "\n"

        # only top-level (non-commented) \input/\include — comments are stripped first
        # so a `% \input{x}` is ignored. We strip comments just for the SCAN, then
        # splice expansions back into the ORIGINAL (comment-bearing) text positions is
        # overkill; simpler: strip comments for the whole flattened blob — pylatexenc
        # also ignores comments, so this is safe and keeps positions sane.
        text = _strip_tex_comments(text)
        return re.sub(r"\\(input|include|subfile)\s*\{([^}]+)\}", _sub, text)

    return expand(main_key, 0)


# ── reference parsing (.bbl preferred, then .bib, then \bibitem in-body) ─────────

_BIB_FIELD_RE = {
    "title":  re.compile(r"title\s*=\s*[{\"]+(.+?)[}\"]+\s*,?\s*$",
                         re.IGNORECASE | re.MULTILINE | re.DOTALL),
    "year":   re.compile(r"year\s*=\s*[{\"]?(\d{4})", re.IGNORECASE),
    "doi":    re.compile(r"doi\s*=\s*[{\"]?([^,}\"]+)", re.IGNORECASE),
    "journal": re.compile(r"(?:journal|booktitle|venue)\s*=\s*[{\"]+(.+?)[}\"]+",
                          re.IGNORECASE | re.DOTALL),
    "author": re.compile(r"author\s*=\s*[{\"]+(.+?)[}\"]+",
                         re.IGNORECASE | re.DOTALL),
}
_ARXIV_IN_TEXT = re.compile(r"ar[Xx]iv[:\s]*(" + _NEW_ID + r"|" + _OLD_ID + r")")
_DOI_IN_TEXT = re.compile(r"10\.\d{4,9}/[-._;()/:A-Za-z0-9]+")


def _clean_field(s: Optional[str]) -> Optional[str]:
    if s is None:
        return None
    s = re.sub(r"[{}]", "", s)
    s = re.sub(r"\s+", " ", s).strip().rstrip(",").strip()
    return s or None


def _parse_bib(text: str) -> list[dict[str, Any]]:
    """Parse a BibTeX .bib into reference dicts. Tolerant brace/field scan (not a full
    BibTeX grammar — good enough for the citation spine: key/authors/title/year/...)."""
    refs = []
    for m in re.finditer(r"@(\w+)\s*\{\s*([^,]+),", text):
        kind, key = m.group(1).lower(), m.group(2).strip()
        if kind in ("comment", "preamble", "string"):
            continue
        # entry body = from this { to the matching } (brace-depth scan)
        start = text.index("{", m.start())
        depth, i = 0, start
        while i < len(text):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    break
            i += 1
        body = text[start:i + 1]
        ent = {
            "key": key,
            "authors": _clean_field(_search1(_BIB_FIELD_RE["author"], body)),
            "title": _clean_field(_search1(_BIB_FIELD_RE["title"], body)),
            "year": _search1(_BIB_FIELD_RE["year"], body),
            "venue": _clean_field(_search1(_BIB_FIELD_RE["journal"], body)),
            "doi": _clean_field(_search1(_BIB_FIELD_RE["doi"], body)),
            "arxiv_id": None,
        }
        am = _ARXIV_IN_TEXT.search(body)
        if am:
            ent["arxiv_id"] = am.group(1)
        elif ent["doi"]:
            # arxiv DOIs (10.48550/arXiv.<id>) carry the id
            dm = re.search(r"arxiv\.(" + _NEW_ID + r"|" + _OLD_ID + r")",
                           ent["doi"], re.IGNORECASE)
            if dm:
                ent["arxiv_id"] = dm.group(1)
        refs.append(ent)
    return refs


def _search1(rx, text) -> Optional[str]:
    m = rx.search(text)
    return m.group(1).strip() if m else None


def _parse_bbl(text: str) -> list[dict[str, Any]]:
    """Parse a compiled .bbl (\\bibitem blocks) into reference dicts. The .bbl is the
    rendered bibliography — author/title aren't separately tagged, so we keep the raw
    item as `title` (best-effort) and harvest any arxiv id / DOI it mentions."""
    refs = []
    # split on \bibitem; each block is one reference. \bibitem[label]{key}
    parts = re.split(r"\\bibitem", text)
    for part in parts[1:]:
        m = re.match(r"\s*(?:\[[^\]]*\])?\s*\{([^}]+)\}", part)
        if not m:
            continue
        key = m.group(1).strip()
        raw = part[m.end():]
        raw = re.split(r"\\bibitem|\\end\s*\{thebibliography\}", raw)[0]
        clean = _to_text(raw)
        ent = {
            "key": key,
            "authors": None,
            "title": clean[:500] or None,
            "year": _search1(re.compile(r"\b(19|20)\d{2}\b"), clean),
            "venue": None,
            "doi": None,
            "arxiv_id": None,
        }
        if ent["year"]:                              # the regex group is only the prefix
            ym = re.search(r"\b((?:19|20)\d{2})\b", clean)
            ent["year"] = ym.group(1) if ym else None
        am = _ARXIV_IN_TEXT.search(raw)
        if am:
            ent["arxiv_id"] = am.group(1)
        dm = _DOI_IN_TEXT.search(raw)
        if dm:
            ent["doi"] = dm.group(0).rstrip(".")
        refs.append(ent)
    return refs


def _extract_references(files: dict[str, bytes],
                        combined_tex: str) -> tuple[list[dict[str, Any]], list[str]]:
    """References, preferring the compiled .bbl, then .bib files, then in-body
    \\bibitem (thebibliography). Returns (refs, warnings)."""
    warnings: list[str] = []
    # 1. compiled .bbl (the rendered bibliography — most reliable structure)
    bbls = [n for n in files if n.lower().endswith(".bbl")]
    if bbls:
        refs = []
        for n in bbls:
            refs += _parse_bbl(_decode(files[n]))
        if refs:
            return refs, warnings
    # 2. .bib source files
    bibs = [n for n in files if n.lower().endswith(".bib")]
    if bibs:
        refs = []
        for n in bibs:
            refs += _parse_bib(_decode(files[n]))
        if refs:
            return refs, warnings
    # 3. in-body \bibitem (thebibliography embedded in the .tex)
    if "\\bibitem" in combined_tex:
        refs = _parse_bbl(combined_tex)
        if refs:
            return refs, warnings
    warnings.append("no references found (.bbl/.bib/\\bibitem all absent)")
    return [], warnings


# ── the pylatexenc structured walk ──────────────────────────────────────────────

def _walk_pylatexenc(combined_tex: str, files: dict[str, bytes],
                     warnings: list[str]) -> dict[str, Any]:
    """Walk the combined LaTeX with pylatexenc, harvesting TYPED formulae / tables /
    figures and a prose body (non-prose envs replaced by a placeholder so the body
    stays readable and the typed atoms aren't duplicated into it)."""
    from pylatexenc.latexwalker import (
        LatexWalker, LatexMathNode, LatexEnvironmentNode, LatexMacroNode,
        LatexGroupNode, LatexCharsNode, LatexCommentNode, LatexSpecialsNode,
    )

    src = combined_tex
    ctx = _pylatexenc_context()
    walker = LatexWalker(src, latex_context=ctx, tolerant_parsing=True)
    nodes, _, _ = walker.get_latex_nodes()

    formulae: list[dict[str, Any]] = []
    tables: list[dict[str, Any]] = []
    figures: list[dict[str, Any]] = []
    body_parts: list[str] = []
    section_stack: list[str] = []

    def cur_section() -> Optional[str]:
        return section_stack[-1] if section_stack else None

    def find_in(env, predicate):
        """Depth-first find of the first descendant node matching predicate."""
        stack = list(getattr(env, "nodelist", []) or [])
        while stack:
            n = stack.pop(0)
            if n is None:
                continue
            if predicate(n):
                return n
            kids = getattr(n, "nodelist", None)
            if kids:
                stack = list(kids) + stack
            argd = getattr(n, "nodeargd", None)
            if argd and argd.argnlist:
                stack = [a for a in argd.argnlist if a is not None] + stack
        return None

    def find_all_in(env, predicate):
        out, stack = [], list(getattr(env, "nodelist", []) or [])
        while stack:
            n = stack.pop(0)
            if n is None:
                continue
            if predicate(n):
                out.append(n)
            kids = getattr(n, "nodelist", None)
            if kids:
                stack = list(kids) + stack
            argd = getattr(n, "nodeargd", None)
            if argd and argd.argnlist:
                stack = [a for a in argd.argnlist if a is not None] + stack
        return out

    def _label_of(env) -> Optional[str]:
        lab = find_in(env, lambda n: isinstance(n, LatexMacroNode)
                      and n.macroname == "label")
        return _macro_arg_text(src, lab) if lab is not None else None

    def _caption_of(env) -> Optional[str]:
        cap = find_in(env, lambda n: isinstance(n, LatexMacroNode)
                      and n.macroname == "caption")
        if cap is None:
            return None
        raw = _macro_arg_text(src, cap)
        return _to_text(raw) if raw else None

    def _table_rows(env):
        """Parse the innermost tabular-family env into rows of cells (split on &, rows
        on \\\\). Returns (rows, tabular_node) or ([], None)."""
        tab = env if env.environmentname in {
            "tabular", "tabular*", "longtable", "tabularx", "supertabular", "array"
        } else find_in(env, lambda n: isinstance(n, LatexEnvironmentNode)
                       and n.environmentname in {
                           "tabular", "tabular*", "longtable", "tabularx",
                           "supertabular", "array"})
        if tab is None:
            return [], None
        body = _env_body_verbatim(src, tab)
        # pylatexenc consumes the {col-spec} arg for `tabular`/`array` but NOT for the
        # less-common variants (longtable / tabularx / supertabular) — for those the
        # leading column-spec (and tabularx's {width}) leaks into the body slice. Strip
        # a leading optional [pos] then up to two {…} groups (width + cols) so the spec
        # never pollutes the first cell.
        if tab.environmentname in {"longtable", "tabularx", "supertabular"}:
            body = re.sub(r"^\s*(?:\[[^\]]*\])?\s*(?:\{[^{}]*\})\s*(?:\{[^{}]*\})?",
                          "", body, count=1)
        rows = []
        for raw_row in re.split(r"\\\\", body):
            raw_row = re.sub(r"\\hline|\\toprule|\\midrule|\\bottomrule|"
                             r"\\cline\s*\{[^}]*\}", "", raw_row)
            if not raw_row.strip():
                continue
            # split on UNescaped & (cell separator)
            cells = re.split(r"(?<!\\)&", raw_row)
            cells = [_to_text(c.strip()) for c in cells]
            if any(c for c in cells):
                rows.append(cells)
        return rows, tab

    def harvest_math_env(env):
        formulae.append({
            "latex": _verbatim(src, env),
            "mathml": None,                          # Phase 2 (latexml/converter)
            "label": _label_of(env),
            "section_context": cur_section(),
            "env": env.environmentname,
        })

    def harvest_table(env):
        rows, _tab = _table_rows(env)
        tables.append({
            "caption": _caption_of(env),
            "label": _label_of(env),
            "rows": rows,
            "latex": _verbatim(src, env),
            "section_context": cur_section(),
        })

    def harvest_figure(env):
        gfx = []
        for g in find_all_in(env, lambda n: isinstance(n, LatexMacroNode)
                             and n.macroname == "includegraphics"):
            tgt = _macro_arg_text(src, g)            # last mandatory arg = file
            if tgt:
                gfx.append(tgt)
        figures.append({
            "caption": _caption_of(env),
            "label": _label_of(env),
            "graphics": gfx,
            "latex": _verbatim(src, env),
            "section_context": cur_section(),
        })

    _SECTION_MACROS = {"section", "subsection", "subsubsection", "section*",
                       "subsection*", "paragraph", "chapter"}

    def walk(node_list, in_document: bool):
        for n in node_list:
            if n is None:
                continue
            if isinstance(n, LatexCommentNode):
                continue
            if isinstance(n, LatexMathNode):
                # inline / display math written with delimiters ($…$, \(…\), \[…\])
                formulae.append({
                    "latex": _verbatim(src, n),
                    "mathml": None,
                    "label": None,
                    "section_context": cur_section(),
                    "env": "inline" if n.displaytype == "inline" else "display",
                })
                continue
            if isinstance(n, LatexEnvironmentNode):
                env = n.environmentname
                if env == "document":
                    walk(n.nodelist, True)
                    continue
                if env in _MATH_ENVS:
                    harvest_math_env(n)
                    continue
                if env in _FIGURE_ENVS:
                    harvest_figure(n)
                    continue
                if env in _TABLE_ENVS or env in ("table", "table*"):
                    harvest_table(n)
                    continue
                if env in ("thebibliography",):
                    continue                         # references handled separately
                if env == "abstract":
                    body_parts.append("\n\n## Abstract\n\n")
                    walk(n.nodelist, in_document)
                    continue
                # generic env (itemize, quote, theorem, …): recurse for prose
                walk(n.nodelist, in_document)
                continue
            if isinstance(n, LatexMacroNode):
                mname = n.macroname
                if mname in _SECTION_MACROS:
                    title = _macro_arg_text(src, n) or ""
                    title = _to_text(title)
                    if title:
                        section_stack[:] = [title]   # flat 1-deep context is enough
                        depth = "#" * (2 if mname.startswith("section")
                                       else 3 if mname.startswith("subsection")
                                       else 4)
                        body_parts.append(f"\n\n{depth} {title}\n\n")
                    continue
                if mname in ("label", "caption", "includegraphics", "cite", "citep",
                             "citet", "ref", "bibliography", "input", "include",
                             "documentclass", "usepackage", "newcommand",
                             "renewcommand", "def", "title", "author", "maketitle",
                             "bibliographystyle", "vspace", "hspace", "noindent"):
                    continue                          # structural / non-prose macros
                # otherwise render the macro (e.g. \textbf{x}) into prose text
                body_parts.append(_to_text(_verbatim(src, n)))
                continue
            if isinstance(n, LatexGroupNode):
                walk(n.nodelist, in_document)
                continue
            if isinstance(n, LatexCharsNode):
                if in_document and n.chars.strip():
                    body_parts.append(n.chars)
                continue
            if isinstance(n, LatexSpecialsNode):
                continue

    walk(nodes, in_document=False)

    body = re.sub(r"[ \t]+", " ", "".join(body_parts))
    body = re.sub(r"\n{3,}", "\n\n", body).strip()
    return {"body": body, "formulae": formulae, "tables": tables, "figures": figures}


def _walk_fallback(combined_tex: str, files: dict[str, bytes],
                   warnings: list[str]) -> dict[str, Any]:
    """Pragmatic env-aware extractor used ONLY when pylatexenc is unavailable. Regex-
    based and therefore weaker on deeply-nested cases, but still env-aware (it pulls
    whole math/table/figure environments verbatim, not blindly). Records a warning so
    the quality drop is visible downstream."""
    warnings.append("pylatexenc unavailable — used the pragmatic regex fallback "
                    "(weaker on nested envs; install pylatexenc for full fidelity)")
    src = _strip_tex_comments(combined_tex)
    formulae, tables, figures = [], [], []

    # display/inline math via delimiters
    for m in re.finditer(r"\\\[(.+?)\\\]|\\\((.+?)\\\)|(?<!\$)\$(?!\$)(.+?)(?<!\\)\$",
                         src, re.DOTALL):
        formulae.append({"latex": m.group(0), "mathml": None, "label": None,
                         "section_context": None, "env": "display" if m.group(1)
                         else "inline"})
    # math / table / figure environments
    def _envs(name):
        return re.finditer(r"\\begin\{" + name + r"\*?\}(.*?)\\end\{" + name + r"\*?\}",
                           src, re.DOTALL)
    for env in ("equation", "align", "gather", "multline", "eqnarray", "alignat",
                "flalign", "displaymath"):
        for m in _envs(env):
            lab = re.search(r"\\label\{([^}]+)\}", m.group(1))
            formulae.append({"latex": m.group(0), "mathml": None,
                             "label": lab.group(1) if lab else None,
                             "section_context": None, "env": env})
    for m in _envs("table"):
        block = m.group(0)
        cap = re.search(r"\\caption\{(.+?)\}", block, re.DOTALL)
        lab = re.search(r"\\label\{([^}]+)\}", block)
        tb = re.search(r"\\begin\{tabular\*?\}(?:\[[^\]]*\])?\{[^}]*\}(.*?)"
                       r"\\end\{tabular\*?\}", block, re.DOTALL)
        rows = []
        if tb:
            for r in re.split(r"\\\\", tb.group(1)):
                r = re.sub(r"\\hline|\\(top|mid|bottom)rule", "", r)
                if r.strip():
                    rows.append([_to_text(c.strip()) for c in re.split(r"(?<!\\)&", r)])
        tables.append({"caption": _to_text(cap.group(1)) if cap else None,
                       "label": lab.group(1) if lab else None, "rows": rows,
                       "latex": block, "section_context": None})
    for m in _envs("figure"):
        block = m.group(0)
        cap = re.search(r"\\caption\{(.+?)\}", block, re.DOTALL)
        lab = re.search(r"\\label\{([^}]+)\}", block)
        gfx = re.findall(r"\\includegraphics(?:\[[^\]]*\])?\{([^}]+)\}", block)
        figures.append({"caption": _to_text(cap.group(1)) if cap else None,
                        "label": lab.group(1) if lab else None, "graphics": gfx,
                        "latex": block, "section_context": None})
    # body: strip the harvested envs + commands to readable text
    body = re.sub(r"\\begin\{(equation|align|gather|multline|eqnarray|table|figure)"
                  r"\*?\}.*?\\end\{\1\*?\}", " ", src, flags=re.DOTALL)
    body = _to_text(body)
    return {"body": body, "formulae": formulae, "tables": tables, "figures": figures}


def extract_eprint(blob_or_dir, card: "ArxivCard" = None) -> dict[str, Any]:
    """Deep TeX extraction (PRIMARY) — Phase 1, BUILT.

    Parse an arxiv TeX e-print into STRUCTURED, TYPED content. The typed
    formulae / tables / figures are CARRIED THROUGH (each its own atom with verbatim
    LaTeX), NEVER flattened into the prose body — preserving math / tables / figures
    as typed content is the entire point.

    `blob_or_dir` accepts a raw e-print tar(.gz)/gzip BLOB (bytes), an already-unpacked
    e-print DIRECTORY path, or a {filename: bytes} MAP (the test/offline form). `card`
    is optional metadata used only to enrich warnings (the parser is card-independent).

    Steps:
      1. assemble the e-print file set; find the main .tex (\\documentclass +
         \\begin{document}); flatten \\input/\\include/\\subfile recursively
         (extension-omitted + flattened-path resolution; cycle-guarded).
      2. walk the combined LaTeX with **pylatexenc** (a real parser — nested envs,
         comments, macro args), harvesting:
           • formulae[]  — every math env ($…$, \\(…\\), \\[…\\], equation/align/
                           gather/multline/eqnarray/…) as {latex VERBATIM, mathml
                           (None — Phase 2), label, section_context, env}.
           • tables[]    — table/tabular → {caption, label, rows (parsed on &/\\\\),
                           latex VERBATIM, section_context}.
           • figures[]   — figure → {caption, label, graphics (\\includegraphics
                           targets), latex VERBATIM, section_context}.
           • body        — readable prose, section structure preserved, commands
                           stripped; the typed envs are NOT duplicated into it.
      3. references[]    — prefer the compiled .bbl, else .bib, else in-body
                           \\bibitem → [{key, authors, title, year, venue, arxiv_id,
                           doi}].

    Returns the deep-extraction output contract:
      {body, formulae[], tables[], figures[], references[], source:"tex", main_tex,
       warnings[]}.

    Robust by construction: pylatexenc handles nested envs / comments / macro defs;
    if pylatexenc is unavailable we fall back to a pragmatic env-aware regex extractor
    (a warning records the fidelity drop) rather than failing the ingest."""
    warnings: list[str] = []
    files = _eprint_file_map(blob_or_dir)
    if not files:
        warnings.append("empty e-print (no files) — nothing to extract")
        return {"body": "", "formulae": [], "tables": [], "figures": [],
                "references": [], "source": "tex", "main_tex": None,
                "warnings": warnings}

    main = _find_main_tex(files)
    if main is None:
        warnings.append("no .tex file in the e-print")
        return {"body": "", "formulae": [], "tables": [], "figures": [],
                "references": [], "source": "tex", "main_tex": None,
                "warnings": warnings}

    combined = _flatten_tex(main, files, warnings)

    try:
        import pylatexenc  # noqa: F401
        result = _walk_pylatexenc(combined, files, warnings)
    except ImportError:
        result = _walk_fallback(combined, files, warnings)
    except Exception as e:                            # noqa: BLE001 — never fail ingest
        warnings.append(f"pylatexenc walk error ({e!r}) — pragmatic fallback")
        result = _walk_fallback(combined, files, warnings)

    refs, ref_warnings = _extract_references(files, combined)
    warnings.extend(ref_warnings)

    result.update({
        "references": refs,
        "source": "tex",
        "main_tex": main,
        "warnings": warnings,
    })
    return result


def extract_pdf(pdf_bytes: bytes, card: "ArxivCard") -> dict[str, Any]:
    """SEAM — PDF extraction (FALLBACK). Phase 2 — NOT built yet.

    Used when the e-print is absent (withdrawn / source-less) or unparseable. Will
    pull the PDF text layer for body, detect tables, rasterise figure regions, and
    parse the reference section. Returns the output contract with source="pdf";
    formulae from PDF are lossy (no LaTeX) and will be flagged low-confidence.

    Phase 2 will add a trivial text extract (pdftotext-class) for the body at minimum;
    until then callers fall back to best_body(card) (title + abstract) when the e-print
    is absent. See GitLab andrea/infrastructure#157 Phase 2."""
    raise NotImplementedError(
        "PDF deep extraction is Phase 2 — see EXTRACTION-MAP + andrea/infrastructure#157")


def best_body(card: "ArxivCard") -> str:
    """Foundation body for ingestion BEFORE the deep extractor exists: the card's
    title + abstract as plain text. This makes an arxiv link immediately ingestable
    and searchable (typed, carded) while the deep TeX/PDF parser is a seam. The deep
    extractor will REPLACE this with the full structured body when built."""
    parts = [card.title.strip(), "", card.abstract.strip()]
    if card.comments:
        parts += ["", f"Comments: {card.comments}"]
    return "\n".join(p for p in parts if p is not None).strip()


# ════════════════════════════════════════════════════════════════════════════════
#  5. TYPED-SOURCE SUBMISSION — card + body → the v2 pipeline as source_type=arxiv
# ════════════════════════════════════════════════════════════════════════════════

def ingest_arxiv(
    arxiv_id: str,
    context_uuid: str = "00000000-0000-0000-0000-000000000000",
    job_id: Optional[str] = None,
    fetch_sources: bool = True,
    body: Optional[str] = None,
) -> dict[str, Any]:
    """Recognise → fetch card → (optionally) stash e-print/PDF + deep-extract → submit.

    The single entry the dissector / web-fetch routes an arxiv link to. Steps:
      1. fetch the official metadata CARD (LookupError if unknown).
      2. if fetch_sources: pull the e-print (TeX) and/or PDF and stash them in the
         store next to the doc, so re-ingests don't re-hit arxiv. (Best-effort — a
         stash failure never blocks the carded ingest.) The fetched e-print bytes are
         returned so we can deep-extract WITHOUT a second fetch.
      3. DEEP EXTRACTION (Phase 1): when the e-print exists, run extract_eprint() to
         get the structured TYPED content — body + formulae[] + tables[] + figures[]
         + references[]. The structured BODY becomes the pipeline content; the typed
         arrays ride under metadata.arxiv.extraction so the chunker/extractor can keep
         them as typed chunks (Phase 3). When the e-print is absent/unparseable we fall
         back to best_body(card) (title + abstract) — extract_pdf is Phase 2.
      4. submit to the v2 pipeline via common/submit.py with source_type="arxiv",
         provider="arxiv", and the full card under metadata.arxiv.

    `body` overrides the body text (skips deep extraction entirely); otherwise the
    deep-extracted body is preferred, then best_body(card) as the foundation fallback.
    """
    card = fetch_card(arxiv_id)
    meta = card.to_metadata()

    # --- stash raw sources beside the doc; keep the e-print bytes for extraction ---
    eprint_bytes: Optional[bytes] = None
    if fetch_sources:
        stash, eprint_bytes = _stash_sources(card, context_uuid)
        if stash:
            meta["arxiv"]["source_stash"] = stash

    # --- deep TeX extraction (Phase 1): prefer the structured body when available ---
    text = body
    if text is None and eprint_bytes:
        try:
            extraction = extract_eprint(eprint_bytes, card)
            if extraction.get("body"):
                text = extraction["body"]
            # carry the TYPED arrays for the chunker/extractor (Phase 3 consumes these);
            # keep the verbatim latex/rows/graphics — DO NOT flatten them into the body.
            meta["arxiv"]["extraction"] = {
                "source": extraction.get("source"),
                "main_tex": extraction.get("main_tex"),
                "formulae": extraction.get("formulae", []),
                "tables": extraction.get("tables", []),
                "figures": extraction.get("figures", []),
                "references": extraction.get("references", []),
                "warnings": extraction.get("warnings", []),
                "counts": {
                    "formulae": len(extraction.get("formulae", [])),
                    "tables": len(extraction.get("tables", [])),
                    "figures": len(extraction.get("figures", [])),
                    "references": len(extraction.get("references", [])),
                },
            }
            logger.info(
                "arxiv deep-extract %s: %d formulae, %d tables, %d figures, %d refs",
                card.arxiv_id,
                len(extraction.get("formulae", [])),
                len(extraction.get("tables", [])),
                len(extraction.get("figures", [])),
                len(extraction.get("references", [])))
        except Exception as e:                          # noqa: BLE001 — never block ingest
            logger.warning("arxiv deep-extract failed for %s: %s — card-body fallback",
                           card.arxiv_id, e)

    if text is None:
        text = best_body(card)

    # Lazy import: common/submit.py pulls boto3 (heavy) — only needed at submit time,
    # keeps recognition/card-fetch import-light for the dissector/web-fetch routers.
    from common.submit import submit as _pipeline_submit
    result = _pipeline_submit(
        content=text,
        title=card.title,
        source_type=SOURCE_TYPE,
        language="en",
        context_uuid=context_uuid,
        metadata=meta,
        job_id=job_id,
    )
    result["arxiv_id"] = card.arxiv_id
    result["arxiv_version"] = card.version
    return result


def _stash_sources(card: "ArxivCard",
                   context_uuid: str) -> tuple[dict[str, Any], Optional[bytes]]:
    """Pull the e-print (TeX) and/or PDF into the store next to a placeholder doc
    path, so re-ingests don't re-hit arxiv. Returns (manifest, eprint_bytes) where the
    manifest is {eprint_key, pdf_key, has_eprint, has_pdf} and eprint_bytes is the raw
    e-print blob (or None) so the caller can deep-extract WITHOUT re-fetching. The
    e-print is fetched even when the store is unavailable, so deep extraction still
    runs. Best-effort: any failure returns a partial manifest and is logged, never
    raised.

    Stash path: pipeline/<ctx>/arxiv/<id>/{eprint.tar, paper.pdf}. The real doc_id is
    assigned by submit.py; this id-keyed path is stable + dedupes re-ingests of the
    same paper. PRIMARY-store write via common.pipeline (savvy-aware)."""
    manifest: dict[str, Any] = {"has_eprint": False, "has_pdf": False}
    id_with_ver = card.arxiv_id + (card.version or "")
    base = f"pipeline/{context_uuid}/arxiv/{id_with_ver}"

    # Fetch the e-print FIRST (independent of store availability) so the deep extractor
    # can run even when the store is down.
    eprint: Optional[bytes] = None
    try:
        eprint = fetch_eprint(id_with_ver)
    except Exception as e:                              # noqa: BLE001
        logger.warning("arxiv stash: e-print fetch failed for %s: %s", id_with_ver, e)

    try:
        from common.pipeline import get_primary_store
        store = get_primary_store()
    except Exception as e:                              # noqa: BLE001
        logger.warning("arxiv stash: store unavailable (%s) — skipping writes", e)
        return manifest, eprint

    if eprint:
        try:
            key = f"{base}/eprint.tar"
            store.put_bytes(key, eprint, content_type="application/x-tar")
            manifest.update(has_eprint=True, eprint_key=key,
                            eprint_bytes=len(eprint))
        except Exception as e:                          # noqa: BLE001
            logger.warning("arxiv stash: e-print store failed for %s: %s",
                           id_with_ver, e)

    # PDF fallback: always stash it too — cheap insurance for the PDF extraction path.
    try:
        pdf = fetch_pdf(id_with_ver)
        if pdf:
            key = f"{base}/paper.pdf"
            store.put_bytes(key, pdf, content_type="application/pdf")
            manifest.update(has_pdf=True, pdf_key=key, pdf_bytes=len(pdf))
    except Exception as e:                              # noqa: BLE001
        logger.warning("arxiv stash: pdf fetch/store failed for %s: %s",
                       id_with_ver, e)
    return manifest, eprint


def route_url(
    url_or_text: str,
    context_uuid: str = "00000000-0000-0000-0000-000000000000",
    job_id: Optional[str] = None,
    **kw,
) -> Optional[dict[str, Any]]:
    """The web-fetch / dissector recognition entry: if `url_or_text` references an
    arxiv id, route it to ingest_arxiv() and return the submit result; else return
    None so the caller falls through to generic webpage handling.

    Wire this at the TOP of any URL-handling path (WebFetch capture, the
    safekeeping-mirror, the curated-list importer):

        res = arxiv.route_url(url, context_uuid=ctx)
        if res is None:
            res = generic_webpage_ingest(url, ctx)   # unchanged fallback
    """
    arxiv_id = recognize(url_or_text)
    if arxiv_id is None:
        return None
    logger.info("arxiv router: %r → id=%s", url_or_text[:80], arxiv_id)
    return ingest_arxiv(arxiv_id, context_uuid=context_uuid, job_id=job_id, **kw)
