#!/usr/bin/env python3
"""Content-stream MCID -> text extraction (file-remediation Stage 2 support).

pdf_table_weave._cell_text() reads /ActualText or /Alt off a table cell's
struct elem — attributes real-world tagged PDFs essentially never set. The
actual visible text lives in the page's content stream, addressed by
marked-content ID (MCID), which the struct tree's /K entries (bare ints or
/MCR dicts) reference. This module parses a page's content stream once and
returns a {mcid: text} map so callers can look up any cell's real text.

Known gap: does NOT descend into Form XObjects (a `Do` operator invoking a
Form XObject with its own nested content stream/resources/marked content is
not followed). Real tagged-table content is essentially always drawn
directly on the page content stream, not inside a Form XObject, so this is
an accepted, documented limitation rather than a correctness bug in
practice.
"""

from __future__ import annotations

import re

import pikepdf

_PAGE_CACHE: list = []  # [(raw_page_obj, {mcid: text}), ...]
_FONT_CACHE: list = []  # [(raw_font_obj, decode_fn), ...]

_ENCODING_CODECS = {
    "/WinAnsiEncoding": "cp1252",
    "/MacRomanEncoding": "mac_roman",
}

_WS_RE = re.compile(r"\s+")
_BFCHAR_BLOCK_RE = re.compile(rb"beginbfchar(.*?)endbfchar", re.S)
_BFRANGE_BLOCK_RE = re.compile(rb"beginbfrange(.*?)endbfrange", re.S)
_HEX_PAIR_RE = re.compile(rb"<([0-9A-Fa-f]+)>\s*<([0-9A-Fa-f]+)>")
_HEX_TRIPLE_RE = re.compile(rb"<([0-9A-Fa-f]+)>\s*<([0-9A-Fa-f]+)>\s*<([0-9A-Fa-f]+)>")
_HEX_RANGE_ARRAY_RE = re.compile(rb"<([0-9A-Fa-f]+)>\s*<([0-9A-Fa-f]+)>\s*\[(.*?)\]", re.S)
_HEX_TOKEN_RE = re.compile(rb"<([0-9A-Fa-f]+)>")


def _cache_get(cache: list, raw):
    """Linear-scan lookup by TRUE object identity (same owning Pdf + same
    objgen). objgen alone is not enough — pikepdf.Object equality (and
    .objgen) compares equal for unrelated objects from two different Pdf
    instances that happen to share an object number, which is the common
    case for small/freshly-generated PDFs (e.g. two single-page documents
    both numbering their page object 4). .same_owner_as() is what actually
    distinguishes the owning Pdf for INDIRECT objects — for direct (never
    made-indirect) objects .same_owner_as() is not a reliable ownership
    signal, so those never match (always a cache miss, safe but uncached).
    A comparison against a since-closed Pdf's stale entry can raise —
    treated as a miss, never propagated."""
    if not getattr(raw, "is_indirect", False):
        return None
    for cached_raw, value in cache:
        try:
            if (getattr(cached_raw, "is_indirect", False)
                    and cached_raw.objgen == raw.objgen
                    and cached_raw.same_owner_as(raw)):
                return value
        except Exception:  # noqa: BLE001
            continue
    return None


def page_mcid_text(page) -> dict:
    """{mcid: text} for every marked-content sequence on `page` that carries
    an /MCID, text collected from Tj/TJ/'/" operators in document order and
    whitespace-normalized. `page` may be a pikepdf.Page or the raw page
    dictionary object (page.obj) — both work with pikepdf.parse_content_stream.
    Never raises: any parse/decode failure yields {} for that page. Cached
    per page (module-level list, matched by true object identity — see
    _cache_get) so the many cells/tables on one page share a single
    content-stream parse."""
    raw = page.obj if hasattr(page, "obj") else page
    cached = _cache_get(_PAGE_CACHE, raw)
    if cached is not None:
        return cached
    try:
        result = _extract(raw)
    except Exception:  # noqa: BLE001 - best-effort report field, must not raise
        result = {}
    _PAGE_CACHE.append((raw, result))
    return result


def _extract(raw) -> dict:
    resources = raw.get("/Resources")
    font_res = resources.get("/Font") if resources and "/Font" in resources else None

    pieces: dict = {}  # mcid -> [str, ...]
    mcid_stack = []  # innermost-open-MCID per nesting depth (or None)
    decode = _decode_latin1

    def emit(text):
        if not mcid_stack or not text:
            return
        mcid = mcid_stack[-1]
        if mcid is None:
            return
        pieces.setdefault(mcid, []).append(text)

    for instr in pikepdf.parse_content_stream(raw):
        op = str(instr.operator)
        operands = instr.operands
        if op == "BDC":
            mcid = _bdc_mcid(operands, resources)
            mcid_stack.append(mcid if mcid is not None else (mcid_stack[-1] if mcid_stack else None))
        elif op == "BMC":
            mcid_stack.append(mcid_stack[-1] if mcid_stack else None)
        elif op == "EMC":
            if mcid_stack:
                mcid_stack.pop()
        elif op == "Tf" and operands and font_res is not None:
            decode = _font_decoder(font_res, operands[0])
        elif op == "Tj" and operands:
            emit(decode(_string_bytes(operands[0])))
        elif op in ("'", '"') and operands:
            emit(decode(_string_bytes(operands[-1])))
        elif op == "TJ" and operands:
            for el in operands[0]:
                if isinstance(el, pikepdf.String):
                    emit(decode(_string_bytes(el)))

    out = {}
    for mcid, chunks in pieces.items():
        text = _WS_RE.sub(" ", "".join(chunks)).strip()
        if text:
            out[mcid] = text
    return out


def _string_bytes(obj) -> bytes:
    try:
        return bytes(obj)
    except Exception:  # noqa: BLE001
        return b""


def _bdc_mcid(operands, resources):
    if len(operands) < 2:
        return None
    props = operands[1]
    if isinstance(props, pikepdf.Name):
        props_res = resources.get("/Properties") if resources and "/Properties" in resources else None
        props = props_res.get(str(props)) if props_res is not None else None
    if props is None or "/MCID" not in props:
        return None
    try:
        return int(props["/MCID"])
    except (TypeError, ValueError):
        return None


def _font_decoder(font_res, name):
    font = font_res.get(str(name)) if font_res is not None else None
    if font is None:
        return _decode_latin1
    cached = _cache_get(_FONT_CACHE, font)
    if cached is not None:
        return cached
    decoder = _build_decoder(font)
    _FONT_CACHE.append((font, decoder))
    return decoder


def _decode_latin1(data: bytes) -> str:
    try:
        return data.decode("latin-1", errors="replace")
    except Exception:  # noqa: BLE001
        return ""


def _build_decoder(font):
    mapping, nbytes = {}, 1
    tounicode = font.get("/ToUnicode")
    if tounicode is not None:
        try:
            parsed_nbytes, parsed_mapping = _parse_tounicode(tounicode.read_bytes())
            if parsed_mapping:
                mapping, nbytes = parsed_mapping, parsed_nbytes
        except Exception:  # noqa: BLE001
            mapping = {}

    codec = _fallback_codec(font)

    def decode(data: bytes) -> str:
        if not data:
            return ""
        if not mapping:
            try:
                return data.decode(codec, errors="replace")
            except LookupError:
                return data.decode("latin-1", errors="replace")
        out = []
        i = 0
        while i + nbytes <= len(data):
            chunk = data[i : i + nbytes]
            if chunk in mapping:
                out.append(mapping[chunk])
            elif nbytes == 1:
                try:
                    out.append(chunk.decode(codec, errors="replace"))
                except LookupError:
                    out.append(chunk.decode("latin-1", errors="replace"))
            i += nbytes
        return "".join(out)

    return decode


def _fallback_codec(font) -> str:
    enc = font.get("/Encoding")
    name = None
    if isinstance(enc, pikepdf.Name):
        name = str(enc)
    elif enc is not None:
        base = enc.get("/BaseEncoding")
        if isinstance(base, pikepdf.Name):
            name = str(base)
    return _ENCODING_CODECS.get(name, "latin-1")


def _hex_to_bytes(hexb: bytes) -> bytes:
    s = hexb.decode("ascii")
    if len(s) % 2:
        s += "0"
    return bytes.fromhex(s)


def _dst_to_str(hexb: bytes) -> str:
    try:
        return _hex_to_bytes(hexb).decode("utf-16-be", errors="replace")
    except Exception:  # noqa: BLE001
        return ""


def _parse_tounicode(data: bytes):
    """Minimal ToUnicode CMap parser: bfchar (1:1) + bfrange (single dst,
    incremented per code, or a bracketed per-code dst array). Returns
    (nbytes, {src_bytes: str}). Best-effort — unparseable/mixed-width
    entries are skipped, never raised."""
    mapping: dict = {}
    nbytes = None

    for block in _BFCHAR_BLOCK_RE.findall(data):
        for src_hex, dst_hex in _HEX_PAIR_RE.findall(block):
            src = _hex_to_bytes(src_hex)
            if nbytes is None:
                nbytes = len(src)
            if len(src) != nbytes:
                continue
            mapping[src] = _dst_to_str(dst_hex)

    for block in _BFRANGE_BLOCK_RE.findall(data):
        for lo_hex, hi_hex, arr_body in _HEX_RANGE_ARRAY_RE.findall(block):
            lo_bytes = _hex_to_bytes(lo_hex)
            width = len(lo_bytes)
            if nbytes is None:
                nbytes = width
            if width != nbytes:
                continue
            lo = int.from_bytes(lo_bytes, "big")
            for offset, dst_hex in enumerate(_HEX_TOKEN_RE.findall(arr_body)):
                code = (lo + offset).to_bytes(width, "big")
                mapping[code] = _dst_to_str(dst_hex)

        for lo_hex, hi_hex, dst_hex in _HEX_TRIPLE_RE.findall(block):
            lo_bytes, hi_bytes = _hex_to_bytes(lo_hex), _hex_to_bytes(hi_hex)
            width = len(lo_bytes)
            if nbytes is None:
                nbytes = width
            if width != nbytes or len(hi_bytes) != width:
                continue
            lo, hi = int.from_bytes(lo_bytes, "big"), int.from_bytes(hi_bytes, "big")
            if hi < lo or hi - lo > 65535:
                continue
            dst_bytes = _hex_to_bytes(dst_hex)
            for code in range(lo, hi + 1):
                src = code.to_bytes(width, "big")
                if len(dst_bytes) == 2:
                    val = int.from_bytes(dst_bytes, "big") + (code - lo)
                    mapping[src] = (val & 0xFFFF).to_bytes(2, "big").decode("utf-16-be", errors="replace")
                else:
                    mapping[src] = dst_bytes.decode("utf-16-be", errors="replace")

    return nbytes or 1, mapping
