#!/usr/bin/env python3
"""In-place PDF font-defect repair (file-remediation Stage 2, Engine E2/E3
library). Ports a reviewer-validated prototype (/tmp/proto_fontfix.py) that
took 5 real STP530 course PDFs from veraPDF 7.21.x failure to COMPLIANT
106/106, pixel-identical.

Research basis (deep-research verified): veraPDF 7.21.7-2 only forbids
U+0000/U+FFFE/U+FEFF in ToUnicode; 7.21.7-1 fails on missing/incomplete
ToUnicode coverage of used glyphs; 7.21.4.2 passes once /CharSet (an
optional FontDescriptor key) is absent. Ghostscript re-distillation was
ruled out — it drops structure trees.

Simple-font (/TrueType, /Type1) passes, in order:
  1. Delete /CharSet from subset-prefixed ([A-Z]{6}+) Type1 FontDescriptors.
  2. Named standard /Encoding (WinAnsi/MacRoman), no /ToUnicode: build a full
     0x00-0xFF ToUnicode CMap via the codec.
  2b. Dict /Encoding with /Differences, no /ToUnicode: build a ToUnicode CMap
      per code — /Differences-named codes resolve via the Adobe Glyph List
      (AGL); every other code falls back to the /BaseEncoding codec decode
      (WinAnsi -> cp1252, MacRoman -> mac_roman) so non-overridden codes stay
      correct, not just present. StandardEncoding / absent /BaseEncoding
      falls back to StandardEncoding glyph names via AGL. veraPDF 7.21.7.
      Reviewer-verified on Pek_2018 (MinionPro/HelveticaNeue/MTEX/
      ZapfDingbats): 6 fonts cleared, fidelity 0.0, body text extracts with
      0 U+FFFD. Blocker fix: an earlier version defaulted ALL /Differences
      bases to StandardEncoding names, silently mismapping ~117 WinAnsi
      codes (e.g. 0xE9 é -> wrong 0xD8 Ø) on every non-overridden code —
      caught by adversarial review, not by any automated gate (ToUnicode
      correctness is invisible to both veraPDF's existence-only check and
      the render-diff fidelity gate).
  3. Fonts with an existing /ToUnicode: rewrite as a full 0-255 coverage CMap,
     preserving existing dst values, gap-filling with U+FFFD, and scrubbing
     forbidden codepoints (U+0000/U+FFFE/U+FEFF -> U+FFFD).

Type0/CID pass:
  4. /Type0 font whose descendant is /CIDFontType2 with no /CIDToGIDMap: set
     /CIDToGIDMap to the name /Identity (ISO 32000 7.21.3.2 requires the
     entry; embedded subset CIDFontType2 fonts use CID==GID). Reviewer-
     verified on Cooper_2018: 4 fonts cleared, fidelity 0.0. Type0 ToUnicode
     (2-byte codespace) is still deferred — see ponytail note below.
  5. Delete /CIDSet from a subset-prefixed CID FontDescriptor — the CID
     counterpart of the /CharSet deletion above (pass 1), same clause
     (7.21.4.2: "CIDSet ... shall identify all glyphs present in the
     embedded font program"). /CIDSet is optional (ISO 32000-1 9.8.2);
     removing it clears the defect without touching the font program.

# ponytail: math fonts (CM*) that falsely declare MacRoman get
# plausible-but-imperfect mappings — this matches pre-existing extraction
# behavior, an honest ceiling, not a regression. Font-PROGRAM defects
# (veraPDF 7.21.4.1, inside the embedded binary) are NOT addressed here;
# that would need fontTools surgery on the embedded font stream, deferred.
# Type0/CID ToUnicode repair (2-byte codespace) is also deferred — only
# /CIDToGIDMap is added for Type0 fonts so far; no failing real-file case
# for CID ToUnicode has been observed yet.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

import pikepdf
from fontTools import agl
from fontTools.encodings.StandardEncoding import StandardEncoding

FORBIDDEN = {0x0000, 0xFFFE, 0xFEFF}
CODEC = {"/MacRomanEncoding": "mac_roman", "/WinAnsiEncoding": "cp1252"}
_SUBSET_RE = re.compile(r"^/[A-Z]{6}\+")


def _scrub(cp):
    if cp in FORBIDDEN or (cp < 0x20 and cp not in (0x09, 0x0A, 0x0D)):
        return 0xFFFD
    return cp


def _parse_existing_cmap(data):
    """bfchar/bfrange -> {code: 'XXXX' hex dst}."""
    txt = data.decode("latin-1")
    m = {}
    mode = None
    for line in txt.split("\n"):
        if "beginbfchar" in line:
            mode = "char"
        elif "beginbfrange" in line:
            mode = "range"
        elif "endbfchar" in line or "endbfrange" in line:
            mode = None
        elif mode:
            hexes = re.findall(r"<([0-9A-Fa-f]+)>", line)
            if mode == "char" and len(hexes) == 2:
                m[int(hexes[0], 16)] = hexes[1].upper()
            elif mode == "range" and len(hexes) == 3 and "[" not in line:
                lo, hi, dst = int(hexes[0], 16), int(hexes[1], 16), int(hexes[2], 16)
                for i, c in enumerate(range(lo, hi + 1)):
                    m[c] = f"{dst + i:04X}"
            elif mode == "range" and "[" in line and len(hexes) >= 3:
                # destination-array form: <lo> <hi> [<d1> <d2> ...] — one dst
                # per code, each possibly multi-unit (ligatures)
                lo, hi = int(hexes[0], 16), int(hexes[1], 16)
                dsts = hexes[2:]
                for i, c in enumerate(range(lo, hi + 1)):
                    if i < len(dsts):
                        m[c] = dsts[i].upper()
    return m


def _cmap_stream(entries):
    """entries: iterable of (code, dst_hex_str) covering 0-255."""
    lines = [
        "/CIDInit /ProcSet findresource begin", "12 dict begin", "begincmap",
        "/CIDSystemInfo <</Registry (Adobe) /Ordering (UCS) /Supplement 0>> def",
        "/CMapName /Adobe-Identity-UCS def", "/CMapType 2 def",
        "1 begincodespacerange", "<00> <FF>", "endcodespacerange",
    ]
    entries = list(entries)
    for i in range(0, len(entries), 100):
        chunk = entries[i:i + 100]
        lines.append(f"{len(chunk)} beginbfchar")
        lines += [f"<{code:02X}> <{dst}>" for code, dst in chunk]
        lines.append("endbfchar")
    lines += ["endcmap", "CMapName currentdict /CMap defineresource pop", "end", "end"]
    return "\n".join(lines).encode("latin-1")


def _base_dst_hex(code, base_encoding):
    """dst hex for one code under the /Differences base encoding, BEFORE any
    /Differences override. WinAnsi/MacRoman must come from the codec decode
    (code -> char -> Unicode) — NOT StandardEncoding glyph names, which
    diverge from WinAnsi at ~117 codes and silently mismap every
    non-overridden accented/punctuation character. StandardEncoding (or no
    /BaseEncoding) falls back to its own glyph-name table via AGL, the
    proven-acceptable default absent a real base encoding."""
    codec = CODEC.get(str(base_encoding)) if base_encoding is not None else None
    if codec is not None:
        try:
            cp = ord(bytes([code]).decode(codec))
        except (ValueError, UnicodeDecodeError):
            cp = 0xFFFD
        return f"{_scrub(cp):04X}"
    name = StandardEncoding[code] if 0 <= code < len(StandardEncoding) else ".notdef"
    return _glyph_name_to_hex(name)


def _differences_overrides(diffs):
    """/Differences array -> {code: glyph_name} override map."""
    overrides = {}
    code = 0
    for item in diffs:
        if isinstance(item, int):
            code = item
            continue
        overrides[code] = str(item).lstrip("/")
        code += 1
    return overrides


def _glyph_name_to_hex(name):
    """Resolve a glyph name to a ToUnicode dst hex string via the AGL.
    Multi-char resolutions (rare) become a multi-unit dst string. Falls back
    to U+FFFD for .notdef or names the AGL can't resolve (veraPDF only
    forbids U+0000/FFFE/FEFF, so FFFD is a legal placeholder)."""
    if not name or name == ".notdef":
        return "FFFD"
    uni = agl.toUnicode(name)
    if not uni:
        return "FFFD"
    return "".join(f"{_scrub(ord(ch)):04X}" for ch in uni)


def _build_tounicode_differences(pdf, font, enc):
    """Pass 2b: build a missing ToUnicode from a dict /Encoding with
    /Differences, covering the font's declared FirstChar..LastChar (the used
    range) — veraPDF 7.21.7. Non-overridden codes fall back to the codec
    decode of /BaseEncoding (WinAnsi/MacRoman) so they're actually correct,
    not just present. Returns True if built."""
    diffs = enc.get("/Differences")
    if diffs is None:
        return False
    base_encoding = enc.get("/BaseEncoding")
    overrides = _differences_overrides(diffs)
    first = int(font.get("/FirstChar", 0))
    last = int(font.get("/LastChar", 255))
    entries = []
    for code in range(first, last + 1):
        if code in overrides:
            entries.append((code, _glyph_name_to_hex(overrides[code])))
        else:
            entries.append((code, _base_dst_hex(code, base_encoding)))
    font["/ToUnicode"] = pdf.make_stream(_cmap_stream(entries))
    return True


def _build_tounicode(pdf, font):
    """Pass 2/2b: build a missing ToUnicode from a named standard encoding,
    or (2b) from a dict /Encoding with /Differences. Returns True if built."""
    enc = font.get("/Encoding")
    if isinstance(enc, pikepdf.Dictionary) and "/Differences" in enc:
        return _build_tounicode_differences(pdf, font, enc)
    codec = CODEC.get(str(enc)) if isinstance(enc, pikepdf.Name) else None
    if codec is None:
        return False
    # Full 0-255 coverage regardless of declared FirstChar/LastChar — veraPDF
    # 7.21.7-1 requires every USED glyph mapped, and declared ranges lie.
    entries = []
    for code in range(0, 256):
        try:
            cp = ord(bytes([code]).decode(codec))
        except (ValueError, UnicodeDecodeError):
            cp = 0xFFFD
        entries.append((code, f"{_scrub(cp):04X}"))
    font["/ToUnicode"] = pdf.make_stream(_cmap_stream(entries))
    return True


def _rewrite_tounicode(pdf, font):
    """Pass 3: rewrite an existing ToUnicode as full 0-255 coverage."""
    tu = font["/ToUnicode"]
    existing = _parse_existing_cmap(tu.read_bytes())
    entries = []
    for code in range(256):
        dst = existing.get(code, "FFFD")
        units = [dst[i:i + 4] for i in range(0, len(dst), 4)]
        units = [f"{_scrub(int(u, 16)):04X}" if len(u) == 4 else u for u in units]
        entries.append((code, "".join(units)))
    font["/ToUnicode"] = pdf.make_stream(_cmap_stream(entries))


def _delete_cidset(font, counts):
    """Pass 5: delete /CIDSet from a subset-prefixed CID FontDescriptor —
    the /CharSet counterpart (pass 1) for Type0/CID fonts, same clause
    (7.21.4.2). Unconditional on subset-prefix, matching the CharSet
    precedent: /CIDSet is optional, so deleting it is always a safe fix for
    a declared-glyph-set mismatch, not just when a mismatch is detected."""
    descendants = font.get("/DescendantFonts")
    if not descendants:
        return
    cid = descendants[0]
    fd = cid.get("/FontDescriptor")
    if fd is None or "/CIDSet" not in fd:
        return
    if not _SUBSET_RE.match(str(font.get("/BaseFont", ""))):
        return
    del fd["/CIDSet"]
    counts["cidsets_deleted"] += 1


def _set_cid_to_gid_map(font, counts):
    """Pass 4: /Type0 + /CIDFontType2 descendant with no /CIDToGIDMap ->
    set /Identity (ISO 32000 7.21.3.2). Reviewer-verified on Cooper_2018."""
    descendants = font.get("/DescendantFonts")
    if not descendants:
        return
    cid = descendants[0]
    if str(cid.get("/Subtype")) != "/CIDFontType2" or "/CIDToGIDMap" in cid:
        return
    # ponytail: assumes CID==GID, true for embedded subset CIDFontType2
    # fonts built by a subsetter (the common case here). The fidelity gate
    # (pixel-diff, exact-zero) is the backstop against a wrong mapping.
    cid["/CIDToGIDMap"] = pikepdf.Name("/Identity")
    counts["cidtogidmap_set"] += 1


def repair_fonts(pdf):
    """Walk every page's /Resources /Font (recursing into Form-XObject
    resources, cycle-guarded), applying all passes in place.
    Returns {"charsets_deleted": n, "cmaps_built": n, "cmaps_rewritten": n,
    "cidtogidmap_set": n, "cidsets_deleted": n}. Differences-encoding
    ToUnicode builds (pass 2b) are folded into "cmaps_built"."""
    counts = {"charsets_deleted": 0, "cmaps_built": 0, "cmaps_rewritten": 0,
              "cidtogidmap_set": 0, "cidsets_deleted": 0}
    seen_fonts = set()
    seen_forms = set()

    def handle_font(font):
        key = font.objgen if font.is_indirect else id(font)
        if key in seen_fonts:
            return
        seen_fonts.add(key)

        subtype = str(font.get("/Subtype"))
        if subtype == "/Type0":
            _set_cid_to_gid_map(font, counts)
            _delete_cidset(font, counts)
            return  # ponytail: Type0 ToUnicode uses a 2-byte codespace (a
            # 0-255 rewrite corrupts it) — deferred until a real failing
            # CID ToUnicode case exists. CIDToGIDMap is the only Type0 fix.
        if subtype not in ("/TrueType", "/Type1"):
            return  # CharSet is Type1-only semantics; other subtypes skip.

        fd = font.get("/FontDescriptor")
        if fd is not None and "/CharSet" in fd and _SUBSET_RE.match(str(font.get("/BaseFont", ""))):
            del fd["/CharSet"]
            counts["charsets_deleted"] += 1
        if "/ToUnicode" not in font:
            if _build_tounicode(pdf, font):
                counts["cmaps_built"] += 1
        else:
            _rewrite_tounicode(pdf, font)
            counts["cmaps_rewritten"] += 1

    def handle_resources(res):
        if res is None:
            return
        for font in (res.get("/Font") or {}).values():
            handle_font(font)
        for xo in (res.get("/XObject") or {}).values():
            if str(xo.get("/Subtype", "")) != "/Form":
                continue
            if xo.is_indirect:
                fkey = xo.objgen
                if fkey in seen_forms:
                    continue
                seen_forms.add(fkey)
            if xo.get("/Resources") is not None:
                handle_resources(xo.get("/Resources"))

    for page in pdf.pages:
        handle_resources(page.get("/Resources"))

    return counts


def main():
    parser = argparse.ArgumentParser(description="In-place PDF font-defect repair")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    try:
        pdf = pikepdf.open(args.input)
        counts = repair_fonts(pdf)
        pdf.save(args.output)
        pdf.close()
    except Exception as exc:
        if args.json:
            print(json.dumps({"error": str(exc)}), file=sys.stderr)
        else:
            print(f"ERROR: {exc}", file=sys.stderr)
        sys.exit(2)

    if args.json:
        print(json.dumps(counts, indent=2))
    else:
        print(f"charsets_deleted={counts['charsets_deleted']} "
              f"cmaps_built={counts['cmaps_built']} cmaps_rewritten={counts['cmaps_rewritten']} "
              f"cidtogidmap_set={counts['cidtogidmap_set']} "
              f"cidsets_deleted={counts['cidsets_deleted']}")


if __name__ == "__main__":
    main()
