#!/usr/bin/env python3
"""papertype-notion.py: save a Notion page, or every page in a Notion database, as a PDF.

Two HTTP calls per page, no dependencies beyond the Python 3.9+ standard library:
  1. GET  https://api.notion.com/v1/pages/<id>/markdown   (Notion returns the page as "enhanced markdown")
  2. POST https://papertype.dev/v1/pdf                     (Markdown -> PDF)
In between, notion_to_markdown() turns Notion's markdown dialect (XML-like tags for callouts, toggles,
columns, mentions; `{color="..."}` attribute lists; single newlines between blocks; tab indentation)
into plain CommonMark/GFM that the renderer understands, and inlines Notion-hosted images as data: URIs
because the renderer fetches nothing from the network.

Usage
  export NOTION_TOKEN=ntn_...          # personal access token or internal-connection token
  export PAPERTYPE_API_KEY=dp_...      # https://papertype.dev/signup
  python3 papertype-notion.py <page URL or id> [out.pdf]
  python3 papertype-notion.py --data-source <data source id> [--filter '<json>'] [--outdir pdfs]

Options
  --theme clean|report|minimal|letter   --page-size A4|Letter|Legal|A5   --toc   --no-images
  --url            print a 24-hour download link instead of saving a file (output=url)
  --properties     for database pages, put a Property | Value table above the content (default on; --no-properties to drop)

MIT License. Copyright (c) 2026 PaperType (papertype.dev).
"""
import argparse, base64, json, os, re, sys, time, urllib.error, urllib.parse, urllib.request

NOTION_API = "https://api.notion.com/v1"
NOTION_VERSION = "2026-03-11"          # required header; see developers.notion.com/reference/versioning
PAPERTYPE_API = os.environ.get("PAPERTYPE_API", "https://papertype.dev/v1/pdf")
IMAGE_BUDGET = 1_400_000               # bytes of base64 per document; the API accepts 2 MB of source per request
IMAGE_MAX = 900_000                    # bytes per image before encoding; bigger images become links


# ----------------------------------------------------------------------------- HTTP
def _http(url, method="GET", headers=None, body=None, retries=3, timeout=60):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method, headers=headers or {})
    if data is not None:
        req.add_header("Content-Type", "application/json")
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.status, dict(r.headers), r.read()
        except urllib.error.HTTPError as e:
            raw = e.read()
            if e.code == 429 and attempt < retries:            # Notion sends Retry-After in seconds
                time.sleep(min(float(e.headers.get("Retry-After", "2") or 2), 60))
                continue
            return e.code, dict(e.headers), raw
    raise RuntimeError("unreachable")


def notion(path, method="GET", body=None, token=None):
    token = token or os.environ.get("NOTION_TOKEN") or sys.exit("NOTION_TOKEN is not set")
    status, headers, raw = _http(NOTION_API + path, method,
                                 {"Authorization": "Bearer " + token, "Notion-Version": NOTION_VERSION}, body)
    try:
        payload = json.loads(raw.decode() or "{}")
    except ValueError:
        payload = {"message": raw[:200].decode(errors="replace")}
    if status >= 400:
        raise RuntimeError(f"Notion {status} on {method} {path}: {payload.get('code', '')} {payload.get('message', '')}".strip())
    return payload


def page_id_from(text):
    """Accept a Notion URL, a 32-hex id or a dashed UUID; return the dashed UUID Notion expects."""
    m = re.findall(r"[0-9a-fA-F]{32}", text.replace("-", ""))
    if not m:
        raise ValueError(f"No 32-character page id found in {text!r}. Copy the page link from Notion (Share -> Copy link).")
    h = m[-1].lower()                                            # the page id is the last one in a URL (view ids come first)
    return f"{h[:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:]}"


# ----------------------------------------------------------------------------- Notion reads
def property_text(prop):
    """One line of text for any page property value (developers.notion.com/reference/page-property-values)."""
    t = prop.get("type")
    v = prop.get(t)
    if v is None:
        return ""
    if t in ("title", "rich_text"):
        return "".join(x.get("plain_text", "") for x in v)
    if t == "number":
        return f"{v:g}" if isinstance(v, (int, float)) else str(v)
    if t in ("select", "status"):
        return v.get("name", "")
    if t == "multi_select":
        return ", ".join(x.get("name", "") for x in v)
    if t == "date":
        return v.get("start", "") + (" → " + v["end"] if v.get("end") else "")
    if t == "checkbox":
        return "Yes" if v else "No"
    if t in ("url", "email", "phone_number"):
        return str(v)
    if t == "people":
        return ", ".join(x.get("name") or x.get("id", "") for x in v)
    if t == "files":
        return ", ".join(x.get("name", "") for x in v)
    if t == "formula":
        return property_text({"type": v.get("type"), v.get("type"): v.get(v.get("type"))})
    if t == "rollup":
        inner = v.get("type")
        if inner == "array":
            return ", ".join(property_text(x) for x in v.get("array", []))
        return property_text({"type": inner, inner: v.get(inner)})
    if t == "relation":
        return f"{len(v)} linked" if isinstance(v, list) else ""
    if t == "unique_id":
        return f"{v.get('prefix') + '-' if v.get('prefix') else ''}{v.get('number', '')}"
    if t in ("created_time", "last_edited_time"):
        return str(v)
    if t in ("created_by", "last_edited_by"):
        return v.get("name") or v.get("id", "")
    return ""


def page_meta(page):
    """(title, {property name: text}) from a page object."""
    title, props = "", {}
    for name, prop in (page.get("properties") or {}).items():
        text = property_text(prop)
        if prop.get("type") == "title":
            title = text
        elif text:
            props[name] = text
    return title or "Untitled", props


def page_markdown(page_id, token=None):
    """Notion's markdown for a page, following unknown_block_ids once when the page was truncated."""
    r = notion(f"/pages/{page_id}/markdown", token=token)
    md = r.get("markdown", "")
    if r.get("truncated"):
        for bid in r.get("unknown_block_ids", []):
            try:
                md += "\n" + notion(f"/pages/{bid}/markdown", token=token).get("markdown", "")
            except RuntimeError:
                pass                                              # 404 = not shared with this token; the <unknown> tag stays
    return md


def query_data_source(ds_id, filter_json=None, token=None):
    """Every page in a data source (Notion paginates at 100; stops at 10,000 per query)."""
    cursor, pages = None, []
    while True:
        body = {"page_size": 100}
        if filter_json:
            body["filter"] = filter_json
        if cursor:
            body["start_cursor"] = cursor
        r = notion(f"/data_sources/{ds_id}/query", "POST", body, token=token)
        pages += [p for p in r.get("results", []) if p.get("object") == "page"]
        if not r.get("has_more"):
            return pages
        cursor = r.get("next_cursor")


# ----------------------------------------------------------------------------- Notion markdown -> CommonMark
ATTRS = re.compile(r'\s*\{(?:[A-Za-z_-]+="[^"]*"\s*)+\}\s*$')
OPEN_TAGS = ("<callout", "<details", "<columns", "<column", "<synced_block", "<synced_block_reference", "<table")
LIST_RE = re.compile(r"^(?:[-*+]|\d+[.)])\s")


def _inline(s):
    """Inline Notion tags -> Markdown/HTML the renderer understands."""
    s = re.sub(r'<mention-(?:user|page|database|data-source|agent)\s+url="([^"]*)"\s*>(.*?)</mention-[a-z-]+>', r"[\2](\1)", s)
    s = re.sub(r'<mention-(?:user|page|database|data-source|agent)\s+url="([^"]*)"\s*/>', r"[\1](\1)", s)
    s = re.sub(r'<mention-date\s+start="([^"]*)"(?:\s+end="([^"]*)")?[^>]*/>', lambda m: m.group(1) + (" → " + m.group(2) if m.group(2) else ""), s)
    s = re.sub(r'<span underline="true">(.*?)</span>', r"<u>\1</u>", s)
    s = re.sub(r'<span color="([a-z]+?)(_bg)?">(.*?)</span>',
               lambda m: f'<span style="{"background:" if m.group(2) else "color:"}{COLORS.get(m.group(1), "inherit")}">{m.group(3)}</span>', s)
    s = re.sub(r'<page url="([^"]*)"[^>]*>(.*?)</page>', r"[\2](\1)", s)
    s = re.sub(r'<database url="([^"]*)"[^>]*>(.*?)</database>', r"[\2](\1)", s)
    s = re.sub(r'<(?:file|pdf|audio|video) src="([^"]*)"[^>]*>(.*?)</(?:file|pdf|audio|video)>', lambda m: f"[{m.group(2) or 'Attachment'}]({m.group(1)})", s)
    s = re.sub(r'<unknown url="([^"]*)"(?:\s+alt="([^"]*)")?\s*/>', lambda m: f"*[{m.group(2) or 'Block not available to this token'}]({m.group(1)})*", s)
    s = re.sub(r"\[\^(https?://[^\]]+)\]", r" ([source](\1))", s)
    s = re.sub(r":([a-z0-9_+-]+):", r"\1", s)                     # custom emoji names
    return s


COLORS = {"gray": "#6b7280", "brown": "#92400e", "orange": "#ea580c", "yellow": "#ca8a04", "green": "#16a34a",
          "blue": "#2563eb", "purple": "#7c3aed", "pink": "#db2777", "red": "#dc2626"}


def notion_to_markdown(md, image_loader=None):
    """Enhanced markdown (Notion's `markdown_version=v1`) -> CommonMark + GFM tables for the renderer.

    Facts this relies on (developers.notion.com, enhanced-markdown and working-with-markdown-content, fetched 2026-10-23):
    one newline between adjacent top-level blocks, tabs for child indentation, `{attr="..."}` lists on a block's first
    line, `<br>` for line breaks inside a block, `<empty-block/>` for an empty line, and the tags handled below.
    Returns (markdown, wants_toc).
    """
    out, stack, wants_toc, fence = [], [], False, None
    budget = IMAGE_BUDGET
    in_table = False
    prev_kind = "blank"

    def push(line, kind):
        nonlocal prev_kind
        # Notion separates blocks with one newline; CommonMark needs a blank line between them, except between
        # consecutive list items and consecutive pipe-table rows. Consecutive callout lines are separate blocks in
        # Notion, so they get a bare ">" between them instead of merging into one paragraph.
        if kind == "quote" and prev_kind == "quote":
            out.append(line[:len(line) - len(line.lstrip("\t"))] + ">")
        elif kind != "blank" and prev_kind not in ("blank", "fence", "table") \
                and not (kind == "list" and prev_kind == "list") and not (kind == "prow" and prev_kind == "prow"):
            out.append("")
        out.append(line)
        prev_kind = kind

    for raw in md.split("\n"):
        indent = len(raw) - len(raw.lstrip("\t"))
        line = raw.lstrip("\t")

        if fence is not None:                                       # inside a code block: verbatim
            out.append(("\t" * max(indent - fence[1], 0)) + line if line else "")
            if line.startswith(fence[0]):
                fence = None
                prev_kind = "fence"
            continue
        m = re.match(r"^(`{3,}|~{3,})", line)
        if m:
            fence = (m.group(1), indent)
            push(("\t" * _depth(stack, indent)) + line, "fence")
            continue

        # Containers that only change nesting: drop the tag, remember how much to dedent.
        while stack and indent < stack[-1][0]:                     # left a container without its closing tag
            stack.pop()
        if in_table:
            if line.startswith("</table>"):
                in_table = False
                out.append("</table>")
                prev_kind = "table"
                continue
            line = re.sub(r'\s+(?:color|width)="[^"]*"', "", line)
            out.append(_inline(line))
            continue
        if line.startswith("<table") and not line.startswith("<table_of_contents"):
            in_table = True
            push(re.sub(r'\s+[a-z-]+="[^"]*"', "", line), "table")
            continue
        if line.startswith("<table_of_contents"):
            wants_toc = True
            continue
        if line.startswith("<empty-block"):
            push("", "blank")
            continue
        if line.startswith("<callout"):
            icon = re.search(r'icon="([^"]*)"', line)
            # Keep text-ish icons (☑ ✓ ★ …); drop emoji from the supplementary planes, which the renderer's fonts lack.
            ic = icon.group(1) if icon and all(ord(ch) <= 0x2BFF for ch in icon.group(1)) else ""
            stack.append((indent, "callout", (ic + " ") if ic else ""))
            continue
        if line.startswith("<details"):
            stack.append((indent, "details", ""))
            continue
        if line.startswith("<summary>"):
            title = re.sub(r"</?summary>", "", line)
            push(("\t" * _depth(stack, indent)) + "**" + _inline(title) + "**", "para")
            continue
        if line.startswith(("<columns", "<column", "<synced_block", "<synced_block_reference")):
            stack.append((indent, "container", ""))
            continue
        if line.startswith("</"):
            if line.startswith("</callout") or line.startswith("</details") or line.startswith("</column") or line.startswith("</synced"):
                if stack:
                    stack.pop()
            continue
        if line.startswith("<custom-block") or line.startswith("<link-preview"):
            continue

        line = ATTRS.sub("", line)
        if not line.strip():
            continue                                                 # Notion strips empty lines; we add our own
        depth = _depth(stack, indent)
        callout = next((c for c in reversed(stack) if c[1] == "callout"), None)

        # Images: inline Notion-hosted files as data: URIs (the renderer never fetches URLs).
        im = re.match(r"^!\[([^\]]*)\]\(([^)\s]+)\)\s*$", line)
        if im and image_loader is not None:
            data, mime = image_loader(im.group(2))
            if data and len(data) <= IMAGE_MAX and budget - len(data) * 4 // 3 > 0:
                b64 = base64.b64encode(data).decode()
                budget -= len(b64)
                line = f"![{im.group(1)}](data:{mime};base64,{b64})"
            else:
                line = f"*[Image{': ' + im.group(1) if im.group(1) else ''}]({im.group(2)})*"
        elif im:
            line = f"*[Image{': ' + im.group(1) if im.group(1) else ''}]({im.group(2)})*"

        line = _inline(line)
        line = re.sub(r"^- \[ \] ", "- ☐ ", line)
        line = re.sub(r"^- \[x\] ", "- ☑ ", line, flags=re.I)
        if line.startswith("$$") and line.strip() == "$$":
            continue
        if callout:
            line = "> " + (callout[2] + line if prev_kind != "quote" else line)
            push(("\t" * depth) + line, "quote")
            continue
        kind = "list" if LIST_RE.match(line) else ("prow" if line.startswith("|") else ("heading" if line.startswith("#") else "para"))
        push(("\t" * depth) + line, kind)

    text = "\n".join(out).strip() + "\n"
    return text, wants_toc


def _depth(stack, indent):
    """How many tabs a line keeps: its own indent minus one per enclosing tag container (whose children are indented)."""
    return max(indent - sum(1 for s in stack if s[0] < indent), 0)


def image_loader(url):
    """Download a Notion-hosted image (signed URL, valid 1 hour) -> (bytes, mime). Returns (None, None) on failure."""
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "papertype-notion/1.0"})
        with urllib.request.urlopen(req, timeout=30) as r:
            mime = (r.headers.get("Content-Type") or "image/png").split(";")[0].strip()
            if not mime.startswith("image/"):
                return None, None
            data = r.read(IMAGE_MAX + 1)
            return (data, mime) if len(data) <= IMAGE_MAX else (None, None)
    except (urllib.error.URLError, OSError, ValueError):
        return None, None


# ----------------------------------------------------------------------------- PaperType
def render_pdf(markdown, *, title, api_key=None, theme="clean", page_size="A4", toc=False, filename="document.pdf", output="pdf"):
    api_key = api_key or os.environ.get("PAPERTYPE_API_KEY") or sys.exit("PAPERTYPE_API_KEY is not set")
    body = {"markdown": markdown, "title": title, "theme": theme, "page_size": page_size, "toc": toc,
            "page_numbers": True, "filename": filename, "output": output}
    status, headers, raw = _http(PAPERTYPE_API, "POST", {"Authorization": "Bearer " + api_key}, body)
    if status != 200:
        try:
            detail = json.loads(raw.decode()).get("detail", raw[:200].decode(errors="replace"))
        except ValueError:
            detail = raw[:200].decode(errors="replace")
        raise RuntimeError(f"PaperType {status}: {detail}")
    if output == "url":
        return json.loads(raw.decode())
    return {"pdf": raw, "pages": headers.get("X-Pages") or headers.get("x-pages")}


def safe_filename(title):
    s = re.sub(r"[^A-Za-z0-9._-]+", "-", title).strip("-.") or "document"
    return s[:80] + ".pdf"


def page_to_pdf(page_id, args, props_table=True):
    page = notion(f"/pages/{page_id}")
    title, props = page_meta(page)
    md, wants_toc = notion_to_markdown(page_markdown(page_id), None if args.no_images else image_loader)
    if not md.lstrip().startswith(f"# {title}\n"):                  # Notion keeps the title out of the content
        md = f"# {title}\n\n" + md
    if props and props_table and args.properties:                 # database pages: Property | Value table under the title
        rows = "\n".join("| %s | %s |" % (k.replace("|", "\\|"), v.replace("|", "\\|").replace("\n", "<br>")) for k, v in props.items())
        first, _, rest = md.lstrip().partition("\n")
        md = f"{first}\n\n| Property | Value |\n|---|---|\n{rows}\n\n{rest.lstrip()}"
    result = render_pdf(md, title=title, theme=args.theme, page_size=args.page_size, toc=args.toc or wants_toc,
                        filename=safe_filename(title), output="url" if args.url else "pdf")
    return title, result


def main(argv=None):
    ap = argparse.ArgumentParser(description="Notion page or database -> PDF via papertype.dev")
    ap.add_argument("target", nargs="?", help="Notion page URL or id")
    ap.add_argument("out", nargs="?", help="output file (default: <title>.pdf)")
    ap.add_argument("--data-source", help="data source id: one PDF per page")
    ap.add_argument("--filter", help="JSON filter object for the data source query")
    ap.add_argument("--outdir", default=".", help="directory for --data-source output")
    ap.add_argument("--theme", default="clean")
    ap.add_argument("--page-size", default="A4")
    ap.add_argument("--toc", action="store_true")
    ap.add_argument("--url", action="store_true", help="print a 24-hour download link instead of saving")
    ap.add_argument("--no-images", action="store_true")
    ap.add_argument("--no-properties", dest="properties", action="store_false")
    args = ap.parse_args(argv)
    if not args.target and not args.data_source:
        ap.error("give a page URL/id or --data-source <id>")

    targets = []
    if args.data_source:
        flt = json.loads(args.filter) if args.filter else None
        pages = query_data_source(args.data_source, flt)
        targets = [p["id"] for p in pages]
        print(f"{len(targets)} page(s) in data source", file=sys.stderr)
    else:
        targets = [page_id_from(args.target)]

    os.makedirs(args.outdir, exist_ok=True)
    for pid in targets:
        title, result = page_to_pdf(pid, args)
        if args.url:
            print(f"{title}\t{result['url']}\t{result['pages']} page(s), expires {result['expires_at']}")
            continue
        path = args.out if (args.out and not args.data_source) else os.path.join(args.outdir, safe_filename(title))
        with open(path, "wb") as f:
            f.write(result["pdf"])
        print(f"{path}\t{result['pages']} page(s)")


if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, ValueError) as e:
        sys.exit(str(e))