#!/usr/bin/env python3
"""
Force PDF text rendering mode to 0 (fill) by rewriting `n Tr` inside BT...ET
for page content streams and Form XObjects recursively.

This targets viewer bugs where text with Tr 4..7 (clipping variants) becomes
invisible unless selected.

Requirements:
  pip install pikepdf

Usage:
  python force_text_render_mode_fill.py in.pdf out.pdf
  python force_text_render_mode_fill.py in.pdf out.pdf --debug
"""

import sys
from typing import List, Set, Tuple
import pikepdf
from pikepdf import Name, Stream, Array, Dictionary

WS = b" \t\r\n\x0c\x00"
DELIMS = b"[]<>()/%"


class Logger:
    def __init__(self, enabled=False):
        self.enabled = enabled
    def log(self, msg):
        if self.enabled:
            print(msg, file=sys.stderr)


def read_literal_string(data: bytes, i: int) -> int:
    n = len(data)
    j = i + 1
    depth = 1
    while j < n and depth > 0:
        c = data[j]
        if c == ord('\\'):
            j += 2
            continue
        if c == ord('('):
            depth += 1
        elif c == ord(')'):
            depth -= 1
        j += 1
    return min(j, n)


def read_hex_string(data: bytes, i: int) -> int:
    n = len(data)
    j = i + 1
    while j < n and data[j] != ord('>'):
        j += 1
    if j < n:
        j += 1
    return j


def read_dict(data: bytes, i: int) -> int:
    n = len(data)
    j = i + 2
    depth = 1
    while j < n and depth > 0:
        c = data[j]
        if c == ord('('):
            j = read_literal_string(data, j)
            continue
        if c == ord('<'):
            if j + 1 < n and data[j + 1] == ord('<'):
                depth += 1
                j += 2
                continue
            j = read_hex_string(data, j)
            continue
        if c == ord('>') and j + 1 < n and data[j + 1] == ord('>'):
            depth -= 1
            j += 2
            continue
        j += 1
    return j


def read_array(data: bytes, i: int) -> int:
    n = len(data)
    j = i + 1
    depth = 1
    while j < n and depth > 0:
        c = data[j]
        if c == ord('('):
            j = read_literal_string(data, j)
            continue
        if c == ord('<'):
            if j + 1 < n and data[j + 1] == ord('<'):
                j = read_dict(data, j)
                continue
            j = read_hex_string(data, j)
            continue
        if c == ord('['):
            depth += 1
        elif c == ord(']'):
            depth -= 1
        j += 1
    return min(j, n)


def read_name(data: bytes, i: int) -> int:
    j = i + 1
    n = len(data)
    while j < n and data[j] not in WS + DELIMS:
        j += 1
    return j


def read_word(data: bytes, i: int) -> int:
    j = i
    n = len(data)
    while j < n and data[j] not in WS + DELIMS:
        j += 1
    if j == i:
        j += 1
    return min(j, n)


def is_number_word(tok: bytes) -> bool:
    try:
        float(tok.decode("latin1"))
        return True
    except Exception:
        return False


def rewrite_force_tr0(data: bytes, logger: Logger, label: str) -> Tuple[bytes, int]:
    out = bytearray()
    i = 0
    n = len(data)

    in_text = False
    # pending numeric operand candidates (byte ranges already emitted)
    recent_words: List[bytes] = []
    tr_changes = 0

    while i < n:
        c = data[i]

        # pass-through tokens that may contain delimiter chars
        if c in WS:
            j = i + 1
            while j < n and data[j] in WS:
                j += 1
            out.extend(data[i:j])
            i = j
            continue

        if c == ord('%'):
            j = i + 1
            while j < n and data[j] not in b"\r\n":
                j += 1
            if j < n and data[j] == ord('\r'):
                j += 1
                if j < n and data[j] == ord('\n'):
                    j += 1
            elif j < n and data[j] == ord('\n'):
                j += 1
            out.extend(data[i:j])
            i = j
            continue

        if c == ord('('):
            j = read_literal_string(data, i)
            out.extend(data[i:j])
            i = j
            continue

        if c == ord('<'):
            if i + 1 < n and data[i + 1] == ord('<'):
                j = read_dict(data, i)
            else:
                j = read_hex_string(data, i)
            out.extend(data[i:j])
            i = j
            continue

        if c == ord('['):
            j = read_array(data, i)
            out.extend(data[i:j])
            i = j
            continue

        if c == ord('/'):
            j = read_name(data, i)
            out.extend(data[i:j])
            i = j
            continue

        j = read_word(data, i)
        tok = data[i:j]

        # Track BT/ET context
        if tok == b"BT":
            in_text = True
            recent_words.clear()
            out.extend(tok)
            i = j
            continue
        if tok == b"ET":
            in_text = False
            recent_words.clear()
            out.extend(tok)
            i = j
            continue

        if in_text and tok == b"Tr":
            # Replace immediate preceding numeric operand with 0 where feasible.
            # We do this by appending canonical " 0 Tr" and relying on consumers;
            # but to avoid duplicate operand confusion, only do this when recent
            # token was a number (common pattern: "4 Tr").
            if recent_words and is_number_word(recent_words[-1]):
                # easiest safe rewrite: emit "0 Tr" and skip emitting current tok
                # BUT previous number already emitted. So instead overwrite logic
                # isn't possible in streaming; canonical workaround is append:
                # "<prevnum> Tr 0 Tr" -> final mode becomes 0.
                out.extend(tok)
                out.extend(b" 0 Tr")
                tr_changes += 1
                logger.log(f"[{label}] forced Tr->0")
                recent_words.clear()
            else:
                out.extend(tok)
            i = j
            continue

        out.extend(tok)

        # track recent word operands only inside text objects
        if in_text:
            if tok and tok not in (b"BT", b"ET"):
                recent_words.append(tok)
                if len(recent_words) > 4:
                    recent_words.pop(0)

        i = j

    return bytes(out), tr_changes


def is_form_xobject(obj) -> bool:
    if not isinstance(obj, Stream):
        return False
    st = obj.get(Name("/Subtype"), None)
    ty = obj.get(Name("/Type"), None)
    return st == Name("/Form") or (ty == Name("/XObject") and st == Name("/Form"))


def page_content_bytes(page) -> bytes:
    c = page.get(Name("/Contents"), None)
    if c is None:
        return b""
    if isinstance(c, Stream):
        return c.read_bytes()
    if isinstance(c, Array):
        parts = []
        for s in c:
            if isinstance(s, Stream):
                parts.append(s.read_bytes())
        return b"\n".join(parts)
    return b""


def set_page_content_bytes(pdf: pikepdf.Pdf, page, data: bytes) -> None:
    page[Name("/Contents")] = pdf.make_stream(data)


def process_form_recursive(form: Stream, visited: Set[Tuple[int, int]], logger: Logger, label: str) -> Tuple[int, int]:
    changed_streams = 0
    tr_changes_total = 0

    objgen = getattr(form, "objgen", None)
    if objgen and objgen in visited:
        return 0, 0
    if objgen:
        visited.add(objgen)

    original = form.read_bytes()
    rewritten, tr_changes = rewrite_force_tr0(original, logger, label)
    if rewritten != original:
        form.write(rewritten)
        changed_streams += 1
    tr_changes_total += tr_changes

    resources = form.get(Name("/Resources"), None)
    if isinstance(resources, Dictionary):
        c, t = process_forms_in_resources(resources, visited, logger, label)
        changed_streams += c
        tr_changes_total += t

    return changed_streams, tr_changes_total


def process_forms_in_resources(resources: Dictionary, visited: Set[Tuple[int, int]], logger: Logger, parent_label: str) -> Tuple[int, int]:
    changed_streams = 0
    tr_changes_total = 0

    xobj = resources.get(Name("/XObject"), None)
    if not isinstance(xobj, Dictionary):
        return 0, 0

    for name_obj, xo in xobj.items():
        if is_form_xobject(xo):
            label = f"{parent_label}/XObject{name_obj}"
            c, t = process_form_recursive(xo, visited, logger, label)
            changed_streams += c
            tr_changes_total += t

    return changed_streams, tr_changes_total


def process_pdf(infile: str, outfile: str, debug: bool) -> None:
    logger = Logger(enabled=debug)
    with pikepdf.open(infile) as pdf:
        modified_streams = 0
        tr_changes_total = 0
        visited_forms: Set[Tuple[int, int]] = set()

        for idx, page in enumerate(pdf.pages, start=1):
            label = f"page[{idx}]"
            original = page_content_bytes(page)
            if original:
                rewritten, tr_changes = rewrite_force_tr0(original, logger, f"{label}/Contents")
                if rewritten != original:
                    set_page_content_bytes(pdf, page, rewritten)
                    modified_streams += 1
                tr_changes_total += tr_changes

            resources = page.get(Name("/Resources"), None)
            if isinstance(resources, Dictionary):
                c, t = process_forms_in_resources(resources, visited_forms, logger, label)
                modified_streams += c
                tr_changes_total += t

        pdf.save(outfile)

    print(f"Wrote {outfile} (modified streams: {modified_streams}, Tr forced to 0: {tr_changes_total})")


def main():
    if len(sys.argv) < 3:
        print("Usage: python force_text_render_mode_fill.py in.pdf out.pdf [--debug]")
        sys.exit(2)

    infile = sys.argv[1]
    outfile = sys.argv[2]
    debug = "--debug" in sys.argv[3:]

    process_pdf(infile, outfile, debug)


if __name__ == "__main__":
    main()