#!/usr/bin/env python3
# PDFVerge IPYNB (Jupyter notebook) to PDF
# Copyright (C) 2026 WebScare (Pvt.) Limited. Licensed under the GNU AGPL-3.0,
# like the rest of PDFVerge's published converters. Source: https://pdfverge.com/source/
"""
    ipynb2pdf.py INPUT.ipynb OUTPUT.pdf [--size letter|a4] [--code show|hide]
                 [--title NAME] [--cpu-secs N]

Lays a Jupyter / Colab / Kaggle notebook out as a PDF with WeasyPrint: markdown
cells (headings, lists, tables, LaTeX math), code cells with syntax colours
and In [n] numbers, and the saved outputs (text, tables, images, errors).
The notebook is NOT run: what you see is what was saved in the file. Nothing
outside the notebook is fetched (no network, no local files) and no script
in it is ever executed.

Result contract (one JSON line on stdout)
  exit 0  {"ok": true, "pages": N, "cells": C, "code": K, "images": I, "math": M, "math_failed": F, "blocked": B}
  exit 2  {"ok": false, "error": "<safe to show the visitor>"}
  exit 3  {"ok": false, "error": "...", "internal": "<log only>"}
"""
import argparse
import base64
import html as H
import json
import logging
import os
import re
import time

_RESULT_FD = os.dup(1)
_devnull = os.open(os.devnull, os.O_WRONLY)
os.dup2(_devnull, 1)
logging.disable(logging.CRITICAL)

PAGE = {"a4": "A4", "letter": "letter"}
ANSI = re.compile(r"\x1b\[[0-9;?]*[A-Za-z]|\x1b\][^\x07]*\x07")


def say(obj, code):
    os.write(_RESULT_FD, (json.dumps(obj) + "\n").encode())
    os._exit(code)


def fail(msg, code=2, internal=None):
    out = {"ok": False, "error": msg}
    if internal:
        out["internal"] = str(internal)[:2000]
    say(out, code)


def limit_resources(mem_mb, cpu_secs):
    try:
        import resource
        if mem_mb:
            b = mem_mb * 1024 * 1024
            resource.setrlimit(resource.RLIMIT_AS, (b, b))
        resource.setrlimit(resource.RLIMIT_CPU, (cpu_secs, cpu_secs + 5))
    except Exception:
        pass


def text_of(v):
    if isinstance(v, list):
        return "".join(str(x) for x in v)
    return str(v or "")


def strip_active(src):
    src = re.sub(r"(?is)<(script|iframe|object|frameset|noscript|template|style)\b.*?</\1\s*>", "", src)
    src = re.sub(r"(?is)<(script|iframe|object|embed|frame|base|link|meta)\b[^>]*>", "", src)
    return re.sub(r"(?is)</?form\b[^>]*>", "", src)


class Stats:
    images = math = math_failed = blocked = 0


# ---------------------------------------------------------------- math
def math_svg(tex, display):
    try:
        import ziamath
        svg = ziamath.Latex(tex.strip(), size=15 if display else 13.5, inline=not display).svg()
        m = re.search(r'viewBox="([-\d.]+) ([-\d.]+) ([-\d.]+) ([-\d.]+)"', svg)
        valign = ""
        if m and not display:
            miny, h = float(m.group(2)), float(m.group(4))
            descent = max(0.0, h + miny)
            valign = "vertical-align:-%.2fpx;" % descent
        uri = "data:image/svg+xml;base64," + base64.b64encode(svg.encode()).decode()
        Stats.math += 1
        if display:
            return '<div class="math-d"><img src="%s" alt="%s"></div>' % (uri, H.escape(tex))
        return '<img class="math-i" style="%s" src="%s" alt="%s">' % (valign, uri, H.escape(tex))
    except Exception:
        Stats.math_failed += 1
        t = H.escape(tex.strip())
        return '<div class="math-d"><code>%s</code></div>' % t if display else "<code>%s</code>" % t


MATH_RE = re.compile(
    r"(?s)(?<!\\)\$\$(.+?)(?<!\\)\$\$"                          # $$ display $$
    r"|\\\[(.+?)\\\]"                                           # \[ display \]
    r"|\\begin\{(equation\*?|align\*?|gather\*?|multline\*?)\}(.+?)\\end\{\3\}"
    r"|(?<![\\$\w])\$(?!\s)([^$\n]+?)(?<![\s\\])\$(?!\w)"         # $inline$ (not $5 and $10)
    r"|\\\((.+?)\\\)")                                          # \( inline \)


def render_markdown(src, attachments):
    """Markdown with math kept out of the markdown parser, then put back as SVG."""
    import markdown
    keep = []

    def hold(html):
        keep.append(html)
        return "\u0000M%d\u0000" % (len(keep) - 1)

    # code spans/blocks first, so $ inside code isn't treated as math
    code_keep = []

    def hold_code(m):
        code_keep.append(m.group(0))
        return "\u0000C%d\u0000" % (len(code_keep) - 1)
    src = re.sub(r"(?s)```.*?```|~~~.*?~~~|`[^`\n]+`", hold_code, src)

    def repl(m):
        if m.group(1) is not None:
            return hold(math_svg(m.group(1), True))
        if m.group(2) is not None:
            return hold(math_svg(m.group(2), True))
        if m.group(4) is not None:
            env = m.group(3)
            body = m.group(4)
            if env.startswith("align"):
                body = "\\begin{aligned}" + body + "\\end{aligned}"
            return hold(math_svg(body, True))
        if m.group(5) is not None:
            return hold(math_svg(m.group(5), False))
        return hold(math_svg(m.group(6), False))
    src = MATH_RE.sub(repl, src)
    src = re.sub("\u0000C(\\d+)\u0000", lambda m: code_keep[int(m.group(1))], src)

    # images attached to the cell: ![x](attachment:name.png)
    def att(m):
        name = m.group(1)
        data = (attachments or {}).get(name) or {}
        for mime in ("image/png", "image/jpeg", "image/gif", "image/svg+xml"):
            if mime in data:
                payload = text_of(data[mime])
                if mime == "image/svg+xml":
                    payload = base64.b64encode(payload.encode()).decode()
                return "data:%s;base64,%s" % (mime, payload.strip())
        return "about:blank"
    src = re.sub(r"attachment:([^\s)\"']+)", att, src)

    out = markdown.markdown(src, extensions=["tables", "fenced_code", "sane_lists"], output_format="html")
    out = strip_active(out)
    return re.sub("\u0000M(\\d+)\u0000", lambda m: keep[int(m.group(1))], out)


# ---------------------------------------------------------------- code
def highlight(code, lang):
    try:
        from pygments import highlight as hl
        from pygments.formatters import HtmlFormatter
        from pygments.lexers import get_lexer_by_name, TextLexer
        try:
            lexer = get_lexer_by_name(lang or "python")
        except Exception:
            lexer = TextLexer()
        return hl(code, lexer, HtmlFormatter(nowrap=True))
    except Exception:
        return H.escape(code)


def pygments_css():
    try:
        from pygments.formatters import HtmlFormatter
        return HtmlFormatter(style="default").get_style_defs(".src")
    except Exception:
        return ""


def render_output(o):
    t = o.get("output_type")
    if t == "stream":
        cls = "out err-stream" if o.get("name") == "stderr" else "out"
        return '<pre class="%s">%s</pre>' % (cls, H.escape(ANSI.sub("", text_of(o.get("text")))))
    if t == "error":
        tb = "\n".join(ANSI.sub("", text_of(x)) for x in (o.get("traceback") or []))
        if not tb.strip():
            tb = "%s: %s" % (o.get("ename", "Error"), o.get("evalue", ""))
        return '<pre class="out error">%s</pre>' % H.escape(tb)
    if t in ("execute_result", "display_data", "pyout", "display"):
        d = o.get("data") or {k: v for k, v in o.items() if k in ("png", "jpeg", "svg", "html", "latex", "text")}
        alias = {"png": "image/png", "jpeg": "image/jpeg", "svg": "image/svg+xml", "html": "text/html", "latex": "text/latex", "text": "text/plain"}
        d = {alias.get(k, k): v for k, v in d.items()}
        for mime in ("image/png", "image/jpeg", "image/gif"):
            if mime in d:
                Stats.images += 1
                return '<div class="out img"><img src="data:%s;base64,%s"></div>' % (mime, re.sub(r"\s+", "", text_of(d[mime])))
        if "image/svg+xml" in d:
            Stats.images += 1
            svg = strip_active(text_of(d["image/svg+xml"]))
            return '<div class="out img"><img src="data:image/svg+xml;base64,%s"></div>' % base64.b64encode(svg.encode()).decode()
        if "text/html" in d:
            h = text_of(d["text/html"])
            stripped = strip_active(h)
            if re.sub(r"(?s)<[^>]+>|\s", "", stripped):
                return '<div class="out html">%s</div>' % stripped
            # only scripts (Plotly, Bokeh, widgets): nothing printable
            plain = text_of(d.get("text/plain", ""))
            return ('<p class="out note">Interactive output (it needs JavaScript, so it can’t be printed).%s</p>'
                    % ((" " + H.escape(plain[:300])) if plain and not plain.startswith("<") else ""))
        if "text/latex" in d:
            tex = text_of(d["text/latex"]).strip().strip("$")
            return '<div class="out">%s</div>' % math_svg(tex, True)
        if "text/markdown" in d:
            return '<div class="out md">%s</div>' % render_markdown(text_of(d["text/markdown"]), {})
        if "text/plain" in d:
            return '<pre class="out">%s</pre>' % H.escape(ANSI.sub("", text_of(d["text/plain"])))
        if any(k.startswith("application/vnd.jupyter.widget") for k in d):
            return '<p class="out note">Interactive widget (not saved as an image, so it can’t be printed).</p>'
    return ""


CSS = """
@page { size: %(size)s; margin: 15mm 14mm 16mm;
  @bottom-left { content: "%(foot)s"; font: 8pt 'DejaVu Sans', sans-serif; color: #7a847f; }
  @bottom-right { content: "Page " counter(page) " of " counter(pages); font: 8pt 'DejaVu Sans', sans-serif; color: #7a847f; } }
html { font: 10pt/1.5 'DejaVu Sans', 'Noto Sans', sans-serif; color: #1b1f24; }
body { margin: 0; }
.nb-head { border-bottom: 1pt solid #d0d7de; padding-bottom: 6pt; margin-bottom: 12pt; }
.nb-head h1 { font-size: 15pt; margin: 0 0 2pt; }
.nb-head p { font-size: 8.5pt; color: #6b7280; margin: 0; }
.cell { margin: 0 0 9pt; }
.md h1 { font-size: 17pt; margin: 12pt 0 6pt; } .md h2 { font-size: 14pt; margin: 11pt 0 5pt; }
.md h3 { font-size: 12pt; margin: 9pt 0 4pt; } .md h4, .md h5, .md h6 { font-size: 10.5pt; margin: 8pt 0 3pt; }
.md h1, .md h2, .md h3 { break-after: avoid; }
.md p { margin: 0 0 6pt; } .md ul, .md ol { margin: 0 0 6pt; padding-left: 18pt; }
.md img { max-width: 100%%; }
.md table, .out.html table { border-collapse: collapse; margin: 4pt 0 8pt; font-size: 8.5pt; }
.md th, .md td, .out.html th, .out.html td { border: 0.6pt solid #d0d7de; padding: 2pt 5pt; text-align: left; }
.out.html td { text-align: right; }
.md th, .out.html th { background: #f3f4f6; }
.md code, .out code { font-family: 'DejaVu Sans Mono', monospace; font-size: 8.8pt; background: #f3f4f6; padding: 0 2pt; border-radius: 2pt; }
.md pre { background: #f6f8fa; padding: 6pt 8pt; border-radius: 3pt; white-space: pre-wrap; font-size: 8.5pt; }
.md pre code { background: none; padding: 0; }
.md blockquote { border-left: 3pt solid #d0d7de; margin: 0 0 6pt; padding-left: 9pt; color: #57606a; }
.code { display: flex; gap: 6pt; }
.prompt { flex: none; width: 13mm; font: 8pt 'DejaVu Sans Mono', monospace; color: #8a8f98; text-align: right; padding-top: 5pt; }
.prompt.outp { color: #c2410c; }
.src { flex: 1; min-width: 0; background: #f6f8fa; border: 0.6pt solid #e1e4e8; border-radius: 3pt; padding: 5pt 7pt;
  font: 8.5pt/1.45 'DejaVu Sans Mono', monospace; white-space: pre-wrap; word-wrap: break-word; margin: 0; }
.outs { flex: 1; min-width: 0; }
.out { margin: 3pt 0 0; font-size: 8.8pt; }
pre.out { font: 8.3pt/1.4 'DejaVu Sans Mono', monospace; white-space: pre-wrap; word-wrap: break-word; padding: 2pt 0; margin: 3pt 0 0; }
pre.out.err-stream { background: #fff5f5; padding: 3pt 6pt; }
pre.out.error { background: #fdecec; color: #8a1c1c; padding: 4pt 7pt; border-radius: 3pt; }
.out.img img { max-width: 100%%; height: auto; break-inside: avoid; }
.out.html { overflow: hidden; } .out.html table { max-width: 100%%; }
.out.note { color: #8a6d1c; font-style: italic; }
.math-d { text-align: center; margin: 6pt 0; } .math-d img { max-width: 100%%; }
.raw { font: 8.5pt 'DejaVu Sans Mono', monospace; white-space: pre-wrap; color: #57606a; }
%(pyg)s
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--size", default="letter")
    ap.add_argument("--code", default="show")
    ap.add_argument("--title", default="")
    ap.add_argument("--cpu-secs", type=int, default=90)
    ap.add_argument("--max-mb", type=int, default=2000)
    a = ap.parse_args()
    limit_resources(a.max_mb, a.cpu_secs)
    t0 = time.time()

    try:
        with open(a.input, "rb") as fh:
            nb = json.loads(fh.read().decode("utf-8-sig"))
    except Exception:
        fail("This file isn't a valid Jupyter notebook (.ipynb). Save it again from Jupyter, Colab or VS Code and retry.")
    if not isinstance(nb, dict):
        fail("This file isn't a valid Jupyter notebook (.ipynb).")
    cells = nb.get("cells")
    if cells is None and isinstance(nb.get("worksheets"), list):     # nbformat 3
        cells = [c for w in nb["worksheets"] for c in (w.get("cells") or [])]
    if not isinstance(cells, list):
        fail("This file isn't a valid Jupyter notebook (.ipynb).")
    meta = nb.get("metadata") or {}
    lang = ((meta.get("language_info") or {}).get("name") or (meta.get("kernelspec") or {}).get("language") or "python")
    kernel = (meta.get("kernelspec") or {}).get("display_name") or ""
    show_code = a.code != "hide"

    title = a.title.strip() or "Notebook"
    body = []
    n_code = 0
    for c in cells:
        if not isinstance(c, dict):
            continue
        ct = c.get("cell_type")
        src = text_of(c.get("source") if "source" in c else c.get("input"))
        if ct == "markdown" or ct == "heading":
            if ct == "heading":
                src = "#" * int(c.get("level") or 1) + " " + src
            if src.strip():
                body.append('<div class="cell md">%s</div>' % render_markdown(src, c.get("attachments")))
        elif ct == "code":
            n_code += 1
            ec = c.get("execution_count", c.get("prompt_number"))
            num = "" if ec is None else str(ec)
            parts = []
            if show_code and src.strip():
                parts.append('<div class="code"><div class="prompt">In [%s]:</div><pre class="src">%s</pre></div>'
                             % (H.escape(num or " "), highlight(src, lang)))
            outs = "".join(render_output(o) for o in (c.get("outputs") or []) if isinstance(o, dict))
            if outs:
                has_result = any(o.get("output_type") in ("execute_result", "pyout") for o in (c.get("outputs") or []) if isinstance(o, dict))
                parts.append('<div class="code"><div class="prompt outp">%s</div><div class="outs">%s</div></div>'
                             % (("Out[%s]:" % H.escape(num)) if has_result and num else "", outs))
            if parts:
                body.append('<div class="cell">%s</div>' % "".join(parts))
        elif ct == "raw":
            if src.strip():
                body.append('<pre class="cell raw">%s</pre>' % H.escape(src))
        if time.time() - t0 > a.cpu_secs:
            fail("This notebook is too large to convert here. Clear some outputs or split it, then try again.")
    if not body:
        fail("This notebook is empty: it has no cells with content.")

    import warnings
    warnings.filterwarnings("ignore")
    try:
        from weasyprint import HTML, CSS as WCSS
        from weasyprint.urls import URLFetcher
    except Exception as ex:
        fail("Notebook conversion is not available right now.", 3, "weasyprint import: " + repr(ex))

    class NbFetcher(URLFetcher):
        """Only data: URIs embedded in the notebook. Never the network or the disk."""
        def fetch(self, url, headers=None):
            if url.lower().startswith("data:"):
                return super().fetch(url, headers)
            if re.match(r"(?i)(https?:)?//", url):
                Stats.blocked += 1
            raise ValueError("blocked")

    fetcher = NbFetcher(allowed_protocols={"data"}, allow_redirects=False, timeout=2)
    foot = (title[:80]).replace("\\", "\\\\").replace('"', '\\"')
    head = '<div class="nb-head"><h1>%s</h1>%s</div>' % (H.escape(title), ('<p>%s</p>' % H.escape(kernel)) if kernel else "")
    doc_html = '<!doctype html><html><head><meta charset="utf-8"><title>%s</title></head><body>%s%s</body></html>' % (
        H.escape(title), head, "\n".join(body))
    css = WCSS(string=CSS % {"size": PAGE.get(a.size, "letter"), "foot": foot, "pyg": pygments_css()})
    try:
        doc = HTML(string=doc_html, url_fetcher=fetcher, base_url="data:/", encoding="utf-8").render(
            stylesheets=[css], optimize_images=True, presentational_hints=True)
        if time.time() - t0 > a.cpu_secs:
            raise TimeoutError()
        doc.metadata.title = title
        doc.write_pdf(a.output)
    except MemoryError:
        fail("This notebook is too large to convert. Clear some outputs and try again.")
    except TimeoutError:
        fail("This notebook took too long to lay out. Clear some outputs or split it, then try again.")
    except Exception as ex:
        fail("This notebook could not be converted.", 3, repr(ex))

    say({"ok": True, "pages": len(doc.pages), "cells": len(cells), "code": n_code, "images": Stats.images,
         "math": Stats.math, "math_failed": Stats.math_failed, "blocked": Stats.blocked,
         "seconds": round(time.time() - t0, 2)}, 0)


if __name__ == "__main__":
    main()