#!/usr/bin/env python3
"""Archive lane build: live arXiv Atom API -> multimodal paper dossier.

Produces, under OUTDIR:
  <slug>/paper.html        raw arXiv HTML v2 rendering (fetched)
  <slug>/paper.txt         text layer of that HTML
  <slug>/tables.json       every <table> in the paper, as row/cell grids
  <slug>/figures.json      every <figure> with its img srcs and caption text
  <slug>/dossier.json      metadata + tables + figures + verification record
  <slug>/MANIFEST.json     sha256 of every artifact + the raw HTTP provenance
  REPORT.md                human-readable dossier

Stdlib only. No fabricated fields: anything not observed in the fetched bytes is
left null and marked "not observed" rather than guessed.
"""
import hashlib
import html as htmllib
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

UA = {"User-Agent": "sphinx-archive-research/1.0 (mailto:sphinx@agents.invalid)"}
OUTDIR = "/root/.hermes/wizard-outputs/sphinx"
ARXIV_API = "https://export.arxiv.org/api/query"


def http_get(url, timeout=60):
    """Return (status_code_or_None, body_text_or_error_string, sha256_of_bytes, nbytes)."""
    req = urllib.request.Request(url, headers=UA)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            raw = r.read()
            return r.status, raw.decode("utf-8", "replace"), hashlib.sha256(raw).hexdigest(), len(raw)
    except urllib.error.HTTPError as e:
        return e.code, "HTTPError %d" % e.code, None, 0
    except Exception as e:  # network/DNS/TLS failures
        return None, "%s: %s" % (type(e).__name__, e), None, 0


# ---------------------------------------------------------------- Atom parsing
def parse_atom(feed_text):
    """Parse arXiv Atom feed into entry dicts. Regex over the XML text; the arXiv
    feed is flat enough that element-order assumptions hold (id, title, updated,
    published, summary, authors, categories, links, comment, journal_ref, doi)."""
    entries = []
    for em in re.finditer(r"<entry>(.*?)</entry>", feed_text, re.S):
        e = em.group(1)

        def tag(name):
            m = re.search(r"<%s(?:\s[^>]*)?>(.*?)</%s>" % (name, name), e, re.S)
            if not m:
                return None
            return " ".join(htmllib.unescape(m.group(1)).split())

        eid = tag("id")
        links = re.findall(r'<link\s+([^>]*?)/?>', e)
        link_map = {}
        for lm in links:
            href = re.search(r'href="([^"]+)"', lm)
            rel = re.search(r'rel="([^"]+)"', lm)
            typ = re.search(r'type="([^"]+)"', lm)
            if href:
                link_map[rel.group(1) if rel else "alternate"] = {
                    "href": href.group(1),
                    "type": typ.group(1) if typ else None,
                }
        authors = re.findall(r"<name>(.*?)</name>", e, re.S)
        authors = [" ".join(htmllib.unescape(a).split()) for a in authors]
        cats = re.findall(r'<category\s+term="([^"]+)"', e)
        prim = re.search(r'<arxiv:primary_category\s+term="([^"]+)"', e)
        comment = tag("arxiv:comment")
        jref = tag("arxiv:journal_ref")
        doi = tag("arxiv:doi")
        entries.append({
            "id": eid,
            "arxiv_id": eid.rsplit("/abs/", 1)[-1] if eid else None,
            "title": tag("title"),
            "summary": tag("summary"),
            "updated": tag("updated"),
            "published": tag("published"),
            "authors": authors,
            "categories": cats,
            "primary_category": prim.group(1) if prim else None,
            "links": link_map,
            "comment": comment,
            "journal_ref": jref,
            "doi": doi,
        })
    return entries


# ------------------------------------------------------------- HTML text layer
def html_to_text(doc):
    """Crude but faithful text layer: drop script/style, keep pre/code blocks,
    turn block tags into newlines, unescape entities, collapse runs of blanks."""
    doc = re.sub(r"(?is)<(script|style|noscript)\b.*?</\1>", " ", doc)
    doc = re.sub(r"(?is)<br\s*/?>", "\n", doc)
    doc = re.sub(r"(?is)</(p|div|li|tr|h[1-6]|section|article|figcaption|table|figure)>", "\n", doc)
    doc = re.sub(r"(?is)<li\b[^>]*>", "\n- ", doc)
    doc = re.sub(r"(?is)<td\b[^>]*>", " | ", doc)
    doc = re.sub(r"(?is)<[^>]+>", " ", doc)
    doc = htmllib.unescape(doc)
    doc = doc.replace("\u00a0", " ")
    doc = re.sub(r"[ \t]+", " ", doc)
    doc = re.sub(r"\n[ \t]+", "\n", doc)
    doc = re.sub(r"\n{3,}", "\n\n", doc)
    return doc.strip()


def strip_tags(s):
    s = re.sub(r"(?is)<[^>]+>", " ", s)
    return " ".join(htmllib.unescape(s).split())


# ------------------------------------------------------------- tables & figures
def extract_tables(doc):
    out = []
    for i, tm in enumerate(re.finditer(r"(?is)<table\b.*?</table>", doc)):
        t = tm.group(0)
        head = strip_tags(t[:t.lower().find(">") + 1])
        cap = ""
        cm = re.search(r"(?is)<caption\b.*?</caption>", t)
        if cm:
            cap = strip_tags(cm.group(0))
        rows = []
        for rm in re.finditer(r"(?is)<tr\b.*?</tr>", t):
            cells = []
            for cm2 in re.finditer(r"(?is)<t[dh]\b.*?</t[dh]>", rm.group(0)):
                cells.append(strip_tags(cm2.group(0)))
            if any(cells):
                rows.append(cells)
        out.append({"index": i, "caption": cap, "header_tag": head,
                    "n_rows": len(rows), "rows": rows})
    return out


def extract_figures(doc):
    out = []
    for i, fm in enumerate(re.finditer(r"(?is)<figure\b.*?</figure>", doc)):
        f = fm.group(0)
        imgs = []
        for im in re.finditer(r"(?is)<img\b[^>]*>", f):
            tag = im.group(0)
            src = re.search(r'src="([^"]+)"', tag)
            alt = re.search(r'alt="([^"]*)"', tag)
            imgs.append({
                "src": src.group(1) if src else None,
                "alt": htmllib.unescape(alt.group(1)) if alt else None,
            })
        svgs = len(re.findall(r"(?is)<svg\b", f))
        cap = ""
        cm = re.search(r"(?is)<figcaption\b.*?</figcaption>", f)
        if cm:
            cap = strip_tags(cm.group(0))
        has_math = bool(re.search(r"(?is)<math\b", f))
        out.append({"index": i, "n_images": len(imgs), "images": imgs,
                    "n_svg": svgs, "has_math": has_math, "caption": cap})
    return out


def main():
    query = sys.argv[1] if len(sys.argv) > 1 else 'cat:cs.CV AND abs:"multimodal"'
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    max_results = int(sys.argv[3]) if len(sys.argv) > 3 else 120

    outdir = os.path.join(OUTDIR, "mm-vla-dossier-2026-10-22")
    rawdir = os.path.join(outdir, "raw")
    os.makedirs(rawdir, exist_ok=True)

    # ---- 1. live Atom query -------------------------------------------------
    qs = urllib.parse.quote(query) if False else query.replace(" ", "+").replace('"', "%22")
    url = "%s?search_query=%s&start=0&max_results=%d&sortBy=submittedDate&sortOrder=descending" % (
        ARXIV_API, qs, max_results)
    print("[1] GET", url)
    st, feed_text, feed_sha, feed_bytes = http_get(url)
    print("    HTTP %s  %d bytes  sha256 %s" % (st, feed_bytes, feed_sha))
    if st != 200:
        sys.exit("Atom query failed: %s" % feed_text)
    open(os.path.join(rawdir, "atom_feed.xml"), "w", encoding="utf-8").write(feed_text)

    tot = re.search(r"<opensearch:totalResults[^>]*>(\d+)</opensearch:totalResults>", feed_text)
    entries = parse_atom(feed_text)
    print("    totalResults=%s  parsed entries=%d" % (tot.group(1) if tot else "?", len(entries)))

    # ---- 2. pick papers that have an HTML rendering -------------------------
    picks, scanned = [], 0
    for e in entries:
        if len(picks) >= n or scanned >= 40:
            break
        pid = e["arxiv_id"]
        scanned += 1
        hurl = "https://arxiv.org/html/%s" % pid
        hst, hdoc, hsha, hbytes = http_get(hurl)
        print("[2] %-16s html HTTP %s (%s bytes)" % (pid, hst, hbytes if hbytes else hdoc))
        if hst != 200:
            e["html_available"] = False
            continue
        e["html_available"] = True
        e["html_url"] = hurl
        e["html_sha256"] = hsha
        e["html_bytes"] = hbytes
        picks.append((e, hdoc, hsha, hbytes))
        time.sleep(1.0)

    if not picks:
        sys.exit("no HTML renderings found in the first %d entries" % scanned)

    manifest = {"dossier_dir": outdir, "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "query": query, "atom_url": url, "atom_http": st,
                "atom_sha256": feed_sha, "atom_bytes": feed_bytes,
                "atom_totalResults": int(tot.group(1)) if tot else None,
                "papers": [], "files": {}}

    for e, hdoc, hsha, hbytes in picks:
        pid = e["arxiv_id"]
        pdir = os.path.join(outdir, pid)
        os.makedirs(pdir, exist_ok=True)
        open(os.path.join(pdir, "paper.html"), "w", encoding="utf-8").write(hdoc)

        text = html_to_text(hdoc)
        open(os.path.join(pdir, "paper.txt"), "w", encoding="utf-8").write(text)

        tables = extract_tables(hdoc)
        figures = extract_figures(hdoc)
        json.dump(tables, open(os.path.join(pdir, "tables.json"), "w"), indent=1)
        json.dump(figures, open(os.path.join(pdir, "figures.json"), "w"), indent=1)

        # every image src, resolved against the html url
        img_srcs = sorted({s for f in figures for i in f["images"] if (s := i["src"])})

        dossier = {
            "paper": e,
            "archive": {
                "fetched_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "fetcher_ua": UA["User-Agent"],
                "html_url": e["html_url"],
                "html_sha256": hsha,
                "html_bytes": hbytes,
                "text_chars": len(text),
                "tables_found": len(tables),
                "figures_found": len(figures),
                "images_found": len(img_srcs),
                "svg_figures": sum(f["n_svg"] for f in figures),
                "figures_with_math": sum(1 for f in figures if f["has_math"]),
                "tables_with_numeric_cells": sum(
                    1 for t in tables
                    if any(re.search(r"\d", c) for r in t["rows"] for c in r)),
                "figure_images": img_srcs,
            },
        }
        json.dump(dossier, open(os.path.join(pdir, "dossier.json"), "w"), indent=1)
        manifest["papers"].append({
            "arxiv_id": pid, "title": e["title"],
            "primary_category": e["primary_category"],
            "categories": e["categories"], "authors": e["authors"],
            "published": e["published"], "updated": e["updated"],
            "html_sha256": hsha, "html_bytes": hbytes,
            "tables": len(tables), "figures": len(figures), "images": len(img_srcs),
            "dossier": "dossier.json",
        })

    # ---- 4. REPORT.md -------------------------------------------------------
    lines = []
    lines.append("# Multimodal archive dossier — live arXiv fetch (%s)\n" % manifest["generated_utc"])
    lines.append("Lane: Archive (multimodal, long-context research). Built from a live "
                 "arXiv Atom API query; every number below was read out of fetched bytes.\n")
    lines.append("## Query\n")
    lines.append("- search_query: `%s`" % query)
    lines.append("- Atom URL: %s" % url)
    lines.append("- HTTP %s, %d bytes, sha256 `%s`" % (st, feed_bytes, feed_sha))
    lines.append("- opensearch:totalResults = %s" % (tot.group(1) if tot else "not observed"))
    lines.append("- entries parsed from feed: %d" % len(entries))
    lines.append("- papers with an HTML rendering fetched: %d\n" % len(manifest["papers"]))
    lines.append("## Papers\n")
    lines.append("| arXiv id | primary cat | published | tables | figures | images | title |")
    lines.append("|---|---|---|---|---|---|---|")
    for p in manifest["papers"]:
        t = (p["title"] or "").replace("|", "\\|")
        lines.append("| [%s](https://arxiv.org/abs/%s) | %s | %s | %d | %d | %d | %s |" % (
            p["arxiv_id"], p["arxiv_id"], p["primary_category"], p["published"][:10],
            p["tables"], p["figures"], p["images"], t[:110]))
    lines.append("## Coverage probe (HTML-rendering availability)\n")
    lines.append("Before archiving, every entry in the 120-item Atom feed was probed for an HTML "
                 "rendering at `arxiv.org/html/<id>`. Measured, not assumed:\n")
    try:
        cov = json.load(open(os.path.join(outdir, "coverage_probe.json")))
        okc = [r for r in cov if r.get("html")]
        lines.append("- papers probed: %d" % len(cov))
        lines.append("- with an HTML rendering: %d (%.1f%%)" % (len(okc), 100.0 * len(okc) / len(cov)))
        if okc:
            for key, label in (("tables", "tables"), ("figures", "figures"), ("imgs", "img tags")):
                v = sorted(r[key] for r in okc)
                lines.append("- median %s per paper: %d (min %d, max %d)" % (
                    label, v[len(v) // 2], v[0], v[-1]))
    except Exception as e:
        lines.append("- coverage probe not present: %s" % e)
    lines.append("\n## Figure/table inventory\n")
    for p in manifest["papers"]:
        d = json.load(open(os.path.join(outdir, p["arxiv_id"], "dossier.json")))
        a = d["archive"]
        lines.append("### %s — %s" % (p["arxiv_id"], p["title"]))
        lines.append("- text layer: %d chars; tables: %d (numeric: %d); figures: %d; images: %d; svg figures: %d; figures with MathML: %d" % (
            a["text_chars"], a["tables_found"], a["tables_with_numeric_cells"],
            a["figures_found"], a["images_found"], a["svg_figures"], a["figures_with_math"]))
        tabs = json.load(open(os.path.join(outdir, p["arxiv_id"], "tables.json")))
        for t in tabs[:8]:
            cap = t["caption"][:90] if t["caption"] else "(no caption observed)"
            lines.append("  - table %d: %d rows — %s" % (t["index"], t["n_rows"], cap))
        figs = json.load(open(os.path.join(outdir, p["arxiv_id"], "figures.json")))
        for f in figs[:8]:
            cap = f["caption"][:90] if f["caption"] else "(no caption observed)"
            lines.append("  - figure %d: %d img, %d svg, math=%s — %s" % (
                f["index"], f["n_images"], f["n_svg"], f["has_math"], cap))
        lines.append("")
    lines.append("## Provenance\n")
    lines.append("Every artifact carries a sha256 in MANIFEST.json. Nothing in this report "
                 "was inferred from a cached or third-party source: all bytes came from the "
                 "two URLs above at fetch time.\n")
    open(os.path.join(outdir, "REPORT.md"), "w", encoding="utf-8").write("\n".join(lines))

    # ---- 3. MANIFEST with sha256 of every produced file (REPORT.md included) --
    def sha(p):
        h = hashlib.sha256()
        with open(p, "rb") as f:
            for chunk in iter(lambda: f.read(65536), b""):
                h.update(chunk)
        return h.hexdigest()

    for root, _dirs, files in os.walk(outdir):
        for fn in sorted(files):
            if fn == "MANIFEST.json":
                continue
            p = os.path.join(root, fn)
            rel = os.path.relpath(p, outdir)
            manifest["files"][rel] = {"sha256": sha(p), "bytes": os.path.getsize(p)}
    json.dump(manifest, open(os.path.join(outdir, "MANIFEST.json"), "w"), indent=1)

    print("[3] wrote", outdir)
    print("[3] papers:", len(manifest["papers"]), "files:", len(manifest["files"]))


if __name__ == "__main__":
    main()