#!/usr/bin/env python3
"""WCAG 1.4.3 text-contrast remediation for PDFs -- background-aware.

v1 of this tool assumed every text span sat on a white page background. That
is wrong for text inside colored callout bars / table-cell fills, and it was
actively harmful: a white label on a failing colored fill was "fixed" by
darkening the *text* to pass against an assumed white page -- which usually
makes real contrast against the *actual* (colored) fill catastrophically
worse (white->grey on a pink bar: 1.3:1). See docs/plan-contrast-bg-aware.md.

v2 detects each span's REAL background (the topmost covering fill rect from
``page.get_drawings()``, or an image -> unfixable, or the page's white if
nothing covers it), skips anything that already passes against that real
background (mandatory -- this is what makes the white->grey regression
impossible), and picks the right side to recolor:
  - white/near-white background, failing       -> darken the TEXT
  - colored fill, text is the DARKER element    -> darken the TEXT (against
    that fill)
  - colored fill, text is the LIGHTER element   -> darken the FILL instead,
    text color is left untouched
  - text sits over an image                     -> unfixable, never guess

Every fix targets an effective ratio with headroom above the raw WCAG
threshold (default floor 5.0, see ``safety_margin``) so boundary-adjacent
fixes (4.51-4.55) that third-party checkers still flag are no longer
possible.

CLI:
    pdf_contrast.py INPUT [-o OUTPUT] [--target-ratio 4.5] [--safety-margin 5.0]
"""
import argparse
import colorsys
import json
import re
import sys
from pathlib import Path

import fitz  # PyMuPDF

sys.path.insert(0, str(Path(__file__).parent))
from wcag_checks import contrast_ratio, relative_luminance

# Default floor for the *effective* target ratio (max(target_ratio,
# safety_margin)). WCAG's own threshold is 4.5; a fix that lands at exactly
# 4.51-4.55 still gets flagged by third-party checkers computing luminance
# with slightly different rounding, so fixes aim for headroom instead.
DEFAULT_SAFETY_MARGIN = 5.0

# Nonstroking fill-color ops: rg (RGB), g (gray), k (CMYK). Uppercase
# (RG/G/K) are stroking colors and are never used for text fill or path
# fill -- skip them.
_COLOR_OP_RE = re.compile(
    rb"(?P<nums>(?:[-+]?[0-9]*\.?[0-9]+\s+){0,3}[-+]?[0-9]*\.?[0-9]+)\s+(?P<op>rg|g|k)\b"
)
_TEXT_SHOW_RE = re.compile(rb"\b(Tj|TJ|'|\")\b")
_PATH_PAINT_RE = re.compile(rb"(?<![A-Za-z0-9/])(f\*?|F|S|s|B\*?|B|b\*?|b|n)(?![A-Za-z0-9])")


def _rgb_to_op_string(rgb):
    r, g, b = (max(0.0, min(1.0, c / 255.0)) for c in rgb)
    return f"{r:.6f} {g:.6f} {b:.6f} rg"


def darken_to_ratio(rgb, target_ratio=4.5, background=(255, 255, 255)):
    """Hue-preserving darken of ``rgb`` until it clears ``target_ratio`` against
    ``background``. Returns ``rgb`` unchanged if it already passes."""
    if contrast_ratio(rgb, background) >= target_ratio:
        return tuple(rgb)
    h, l, s = colorsys.rgb_to_hls(*(c / 255.0 for c in rgb))
    lo, hi = 0.0, l
    best = (0, 0, 0)
    for _ in range(40):
        mid = (lo + hi) / 2
        candidate = tuple(round(c * 255) for c in colorsys.hls_to_rgb(h, mid, s))
        if contrast_ratio(candidate, background) >= target_ratio:
            lo = mid
            best = candidate
        else:
            hi = mid
    return best


def _darken_multi(rgb, others, target_ratio):
    """Hue-preserving darken of ``rgb`` until it clears ``target_ratio``
    against EVERY color in ``others`` simultaneously (used when the same
    original color must satisfy more than one distinct real background, or
    when one fill must satisfy more than one overlaying text color). Returns
    ``None`` if even the fully-darkened value can't satisfy all of them at
    once -- e.g. genuinely conflicting light/dark constraints -- rather than
    faking a color that doesn't actually pass."""
    others = list(others)

    def passes(candidate):
        return all(contrast_ratio(candidate, o) >= target_ratio for o in others)

    if passes(tuple(rgb)):
        return tuple(rgb)

    h, l, s = colorsys.rgb_to_hls(*(c / 255.0 for c in rgb))
    lo, hi = 0.0, l
    best = None
    for _ in range(40):
        mid = (lo + hi) / 2
        candidate = tuple(round(c * 255) for c in colorsys.hls_to_rgb(h, mid, s))
        if passes(candidate):
            best = candidate
            lo = mid
        else:
            hi = mid
    if best is not None:
        return best
    return (0, 0, 0) if passes((0, 0, 0)) else None


def _mask_paren_strings(data):
    """Blank the interior of PDF literal strings ``(...)`` (balanced nested
    parens, backslash escapes respected) so a lone letter inside a kerned TJ
    fragment -- e.g. the trailing ``f`` in ``(Incorporate f)`` -- can never be
    mistaken for a path-paint operator token (``f``/``F``/``S``/``s``/``B``/
    ``b``/``n``) by the classifier's lookahead regexes below. Length is
    preserved so match offsets stay comparable."""
    out = bytearray(data)
    i, n, depth = 0, len(out), 0
    while i < n:
        c = out[i]
        if depth == 0 and c == 0x28:  # '('
            depth = 1
            out[i] = 0x20
            i += 1
            continue
        if depth > 0:
            if c == 0x5C and i + 1 < n:  # backslash escape: blank both bytes
                out[i] = 0x20
                out[i + 1] = 0x20
                i += 2
                continue
            if c == 0x28:
                depth += 1
            elif c == 0x29:
                depth -= 1
            out[i] = 0x20
            i += 1
            continue
        i += 1
    return bytes(out)


def _classify_occurrence(content, match_end):
    """Look ahead from a color-op match end to the next painting operator.
    Returns 'text', 'path', or 'unknown' (no recognizable paint op found
    before another color op / end of stream)."""
    window = content[match_end : match_end + 400]
    masked = _mask_paren_strings(window)
    text_m = _TEXT_SHOW_RE.search(masked)
    path_m = _PATH_PAINT_RE.search(masked)
    if text_m and (not path_m or text_m.start() < path_m.start()):
        return "text"
    if path_m and (not text_m or path_m.start() < text_m.start()):
        return "path"
    return "unknown"


def _clamp255(v):
    return max(0, min(255, int(round(v))))


def _cmyk_to_rgb_fitz(c, m, y, k):
    """Convert CMYK (0..1 floats) to RGB using PyMuPDF's own color-management
    pipeline, so this matches what ``page.get_text("dict")`` reports for the
    same operands. A naive ``255*(1-c)*(1-k)`` formula diverges from fitz's
    conversion enough that failing CMYK colors never match the detected RGB
    list (see PR fix/pdf-contrast-cmyk)."""
    samples = bytes(_clamp255(v * 255) for v in (c, m, y, k))
    pix = fitz.Pixmap(fitz.csCMYK, 1, 1, samples, 0)
    rgb_pix = fitz.Pixmap(fitz.csRGB, pix)
    return rgb_pix.pixel(0, 0)


def _op_to_rgb(nums_str, op):
    parts = [float(x) for x in nums_str.split()]
    if op in (b"rg", "rg"):
        if len(parts) != 3:
            return None
        r, g, b = parts
        return (round(r * 255), round(g * 255), round(b * 255))
    if op in (b"g", "g"):
        if len(parts) != 1:
            return None
        v = parts[0]
        gray = round(v * 255)
        return (gray, gray, gray)
    if op in (b"k", "k"):
        if len(parts) != 4:
            return None
        c, m, y, k = parts
        return _cmyk_to_rgb_fitz(c, m, y, k)
    return None


def _page_fill_rects(page):
    """Filled-path drawings on ``page`` as ``(rect, rgb0_255, seqno)``,
    sorted ascending by ``seqno`` (paint order) so the LAST match for a given
    point is the topmost one. Near-transparent fills (opacity < 0.5) are
    skipped -- they don't read as a real visual background."""
    fills = []
    for d in page.get_drawings():
        fill = d.get("fill")
        rect = d.get("rect")
        if fill is None or rect is None:
            continue
        opacity = d.get("fill_opacity")
        if opacity is not None and opacity < 0.5:
            continue
        rgb = tuple(_clamp255(c * 255) for c in fill)
        fills.append((rect, rgb, d.get("seqno", 0)))
    fills.sort(key=lambda t: t[2])
    return fills


def _page_image_rects(page):
    """Placement rects of every image on ``page``, across all xrefs."""
    rects = []
    for img in page.get_images(full=True):
        try:
            rects.extend(page.get_image_rects(img[0]))
        except Exception:
            continue
    return rects


def _detect_background(cx, cy, fills, image_rects, default_bg):
    """Real background of point ``(cx, cy)``: the topmost covering fill rect,
    else (if instead covered by an image) a sentinel meaning "can't recolor",
    else ``default_bg`` (the page's own background, normally white). Returns
    ``("fill", rgb)`` | ``("image", None)`` | ``("white", default_bg)``."""
    point = fitz.Point(cx, cy)
    topmost = None
    for rect, rgb, _seqno in fills:  # ascending seqno -> last match = topmost
        if rect.contains(point):
            topmost = rgb
    if topmost is not None:
        return ("fill", topmost)
    for rect in image_rects:
        if rect.contains(point):
            return ("image", None)
    return ("white", default_bg)


def _scan_page_spans(page):
    """Non-whitespace text spans on ``page``: text, fill color, and bbox
    center, for background detection."""
    spans = []
    d = page.get_text("dict")
    for block in d.get("blocks", []):
        for line in block.get("lines", []):
            for span in line.get("spans", []):
                text = span.get("text", "")
                if text.strip() == "":
                    continue
                color_int = span.get("color", 0)
                rgb = ((color_int >> 16) & 255, (color_int >> 8) & 255, color_int & 255)
                x0, y0, x1, y1 = span.get("bbox", (0, 0, 0, 0))
                spans.append({"text": text, "rgb": rgb, "cx": (x0 + x1) / 2, "cy": (y0 + y1) / 2})
    return spans


def remediate_contrast(input_path, output_path=None, target_ratio=4.5,
                        background=(255, 255, 255), safety_margin=DEFAULT_SAFETY_MARGIN):
    """Recolor low-contrast text -- or, for light text on a colored fill that
    fails, the FILL itself -- in a PDF, measured against each span's real
    background. Returns result dict with ``fixed`` (text recolors),
    ``fills_darkened`` (fill recolors), and ``unfixable`` (with reasons)."""
    if output_path is None:
        output_path = str(Path(input_path).with_suffix("")) + ".contrast_fixed.pdf"

    effective_target = max(target_ratio, safety_margin)

    result = {
        "ok": False,
        "target_ratio": target_ratio,
        "safety_margin": safety_margin,
        "effective_target_ratio": effective_target,
        "fixed": [],
        "fills_darkened": [],
        "unfixable": [],
        "unchanged_black_text": 0,
        "already_passing": 0,
    }

    doc = fitz.open(input_path)

    # 1. Per page: detect the REAL background of every non-whitespace span
    # and bucket it. Anything already passing against its real background is
    # left alone -- mandatory, this is what makes the white-on-fill ->
    # grey-on-fill regression impossible.
    text_groups = {}  # (pno, text_rgb) -> {"backgrounds": set, "sample_text": str, "worst_ratio": float}
    fill_groups = {}  # (pno, fill_rgb) -> {"texts": set, "sample_text": str, "worst_ratio": float}
    declared_unfixable_keys = set()  # (pno, rgb) already reported unfixable

    for pno, page in enumerate(doc):
        fills = _page_fill_rects(page)
        image_rects = _page_image_rects(page)
        for span in _scan_page_spans(page):
            rgb = span["rgb"]
            kind, bg_rgb = _detect_background(span["cx"], span["cy"], fills, image_rects, background)

            if kind == "image":
                key = (pno, rgb)
                if key not in declared_unfixable_keys:
                    declared_unfixable_keys.add(key)
                    result["unfixable"].append({
                        "page": pno, "text": span["text"][:60], "rgb": list(rgb),
                        "ratio": None,
                        "reason": "text over image background — cannot recolor",
                    })
                continue

            ratio = contrast_ratio(rgb, bg_rgb)
            if ratio >= effective_target:
                result["already_passing"] += 1
                if rgb == (0, 0, 0):
                    result["unchanged_black_text"] += 1
                continue

            if kind == "fill" and relative_luminance(rgb) > relative_luminance(bg_rgb):
                # Text is the LIGHTER element on a colored fill -> darken the
                # fill, never the text (this is the defect-2 regression path).
                grp = fill_groups.setdefault(
                    (pno, bg_rgb),
                    {"texts": set(), "sample_text": span["text"][:60], "worst_ratio": ratio},
                )
                grp["texts"].add(rgb)
                grp["worst_ratio"] = min(grp["worst_ratio"], ratio)
            else:
                # White background, or a colored fill where text is already
                # the darker element -> darken the text against that bg.
                grp = text_groups.setdefault(
                    (pno, rgb),
                    {"backgrounds": set(), "sample_text": span["text"][:60], "worst_ratio": ratio},
                )
                grp["backgrounds"].add(bg_rgb)
                grp["worst_ratio"] = min(grp["worst_ratio"], ratio)

    # 2. Resolve each group to a single new color, or mark unreachable.
    text_fix_by_page = {}   # pno -> {rgb: new_rgb}
    text_info_by_page = {}  # pno -> {rgb: {sample_text, ratio_before, backgrounds}}
    for (pno, rgb), grp in text_groups.items():
        backgrounds = list(grp["backgrounds"])
        if len(backgrounds) == 1:
            new_rgb = darken_to_ratio(rgb, effective_target, backgrounds[0])
            reached = new_rgb is not None and contrast_ratio(new_rgb, backgrounds[0]) >= effective_target
        else:
            new_rgb = _darken_multi(rgb, backgrounds, effective_target)
            reached = new_rgb is not None
        if not reached:
            key = (pno, rgb)
            if key not in declared_unfixable_keys:
                declared_unfixable_keys.add(key)
                result["unfixable"].append({
                    "page": pno, "text": grp["sample_text"], "rgb": list(rgb),
                    "ratio": round(grp["worst_ratio"], 3),
                    "reason": "cannot reach target contrast by darkening alone against all of its "
                    f"real backgrounds ({[list(b) for b in backgrounds]})",
                })
            continue
        text_fix_by_page.setdefault(pno, {})[rgb] = new_rgb
        text_info_by_page.setdefault(pno, {})[rgb] = {
            "sample_text": grp["sample_text"], "ratio_before": grp["worst_ratio"], "backgrounds": backgrounds,
        }

    fill_fix_by_page = {}   # pno -> {fill_rgb: new_fill_rgb}
    fill_info_by_page = {}  # pno -> {fill_rgb: {sample_text, ratio_before, texts}}
    for (pno, fill_rgb), grp in fill_groups.items():
        texts = list(grp["texts"])
        if len(texts) == 1:
            new_fill = darken_to_ratio(fill_rgb, effective_target, texts[0])
            reached = new_fill is not None and contrast_ratio(texts[0], new_fill) >= effective_target
        else:
            new_fill = _darken_multi(fill_rgb, texts, effective_target)
            reached = new_fill is not None
        if not reached:
            result["unfixable"].append({
                "page": pno, "text": grp["sample_text"], "rgb": list(fill_rgb), "kind": "fill",
                "ratio": round(grp["worst_ratio"], 3),
                "reason": "fill color cannot be darkened enough to pass against the overlaying "
                f"text color(s) {[list(t) for t in texts]}",
            })
            # The overlaying text is left unchanged and stays failing against
            # the real (un-darkened) fill -- declare it too, honestly.
            for t in texts:
                key = (pno, t)
                if key not in declared_unfixable_keys:
                    declared_unfixable_keys.add(key)
                    result["unfixable"].append({
                        "page": pno, "text": grp["sample_text"], "rgb": list(t),
                        "ratio": round(grp["worst_ratio"], 3),
                        "reason": f"sits on fill {list(fill_rgb)} which cannot be darkened enough to pass",
                    })
            continue
        fill_fix_by_page.setdefault(pno, {})[fill_rgb] = new_fill
        fill_info_by_page.setdefault(pno, {})[fill_rgb] = {
            "sample_text": grp["sample_text"], "ratio_before": grp["worst_ratio"], "texts": texts,
        }

    # 3. Rewrite content streams in one pass per stream: a text-show
    # occurrence matching a text target gets the new TEXT color; a
    # path-paint occurrence matching a fill target gets the new FILL color.
    # Never the other way -- that's the whole point of the classifier.
    rewritten_text_by_page = {}
    rewritten_fill_by_page = {}
    for pno, page in enumerate(doc):
        text_targets = text_fix_by_page.get(pno, {})
        fill_targets = fill_fix_by_page.get(pno, {})
        if not text_targets and not fill_targets:
            continue
        rewritten_text_by_page[pno] = set()
        rewritten_fill_by_page[pno] = set()
        for xref in page.get_contents():
            raw = doc.xref_stream(xref)
            if raw is None:
                continue

            out = bytearray()
            last = 0
            for m in _COLOR_OP_RE.finditer(raw):
                op = m.group("op")
                if op in (b"RG", b"G", b"K"):
                    continue  # stroking, never a text or path fill
                rgb = _op_to_rgb(m.group("nums"), op)
                if rgb is None:
                    continue
                kind = _classify_occurrence(raw, m.end())

                new_rgb = None
                if kind == "text":
                    for failing_rgb, replacement in text_targets.items():
                        if all(abs(rgb[i] - failing_rgb[i]) <= 2 for i in range(3)):
                            new_rgb = replacement
                            rewritten_text_by_page[pno].add(failing_rgb)
                            break
                elif kind == "path":
                    for failing_rgb, replacement in fill_targets.items():
                        if all(abs(rgb[i] - failing_rgb[i]) <= 2 for i in range(3)):
                            new_rgb = replacement
                            rewritten_fill_by_page[pno].add(failing_rgb)
                            break

                if new_rgb is None:
                    continue
                out += raw[last : m.start()]
                out += _rgb_to_op_string(new_rgb).encode("latin-1")
                last = m.end()
            out += raw[last:]
            new_bytes = bytes(out)
            if new_bytes != raw:
                doc.update_stream(xref, new_bytes)

    # 4. Record fixed / unfixable based on what actually got rewritten.
    for pno, targets in text_fix_by_page.items():
        rewritten = rewritten_text_by_page.get(pno, set())
        infos = text_info_by_page.get(pno, {})
        for rgb, new_rgb in targets.items():
            info = infos.get(rgb, {})
            if rgb in rewritten:
                backgrounds = info.get("backgrounds", [(255, 255, 255)])
                result["fixed"].append({
                    "page": pno, "text": info.get("sample_text", ""),
                    "from_rgb": list(rgb), "to_rgb": list(new_rgb),
                    "ratio_before": round(info.get("ratio_before", 0), 3),
                    "ratio_after": round(min(contrast_ratio(new_rgb, bg) for bg in backgrounds), 3),
                })
            else:
                key = (pno, rgb)
                if key not in declared_unfixable_keys:
                    declared_unfixable_keys.add(key)
                    result["unfixable"].append({
                        "page": pno, "text": info.get("sample_text", ""),
                        "rgb": list(rgb), "ratio": round(info.get("ratio_before", 0), 3),
                        "reason": "color operator not found as a text-show fill (scn/pattern/ICC "
                        "separation, or precedes a path paint op rather than a text-show op)",
                    })

    fills_darkened_map = {}  # (from_rgb, to_rgb) -> {"pages": set, "ratio_after": float}
    for pno, targets in fill_fix_by_page.items():
        rewritten = rewritten_fill_by_page.get(pno, set())
        infos = fill_info_by_page.get(pno, {})
        for fill_rgb, new_fill in targets.items():
            info = infos.get(fill_rgb, {})
            texts = info.get("texts", [(255, 255, 255)])
            if fill_rgb in rewritten:
                key = (fill_rgb, new_fill)
                entry = fills_darkened_map.setdefault(key, {"pages": set(), "ratio_after": None})
                entry["pages"].add(pno)
                ratio_after = min(contrast_ratio(t, new_fill) for t in texts)
                entry["ratio_after"] = ratio_after if entry["ratio_after"] is None else min(entry["ratio_after"], ratio_after)
            else:
                result["unfixable"].append({
                    "page": pno, "text": info.get("sample_text", ""), "rgb": list(fill_rgb), "kind": "fill",
                    "ratio": round(info.get("ratio_before", 0), 3),
                    "reason": "fill color operator not found preceding a path-paint op (pattern/ICC "
                    "separation, or precedes a text-show op instead)",
                })
                for t in texts:
                    key = (pno, t)
                    if key not in declared_unfixable_keys:
                        declared_unfixable_keys.add(key)
                        result["unfixable"].append({
                            "page": pno, "text": info.get("sample_text", ""), "rgb": list(t),
                            "ratio": round(info.get("ratio_before", 0), 3),
                            "reason": "sits on a fill whose color operator could not be rewritten",
                        })

    for (from_rgb, to_rgb), entry in fills_darkened_map.items():
        result["fills_darkened"].append({
            "from_rgb": list(from_rgb), "to_rgb": list(to_rgb),
            "pages": sorted(entry["pages"]),
            "ratio_after": round(entry["ratio_after"], 3) if entry["ratio_after"] is not None else None,
        })

    doc.save(output_path)
    doc.close()

    # 5. Verify: re-open the OUTPUT, re-detect REAL backgrounds (never
    # assume white), re-probe. Anything still failing that wasn't already
    # declared unfixable is a genuine regression.
    verify_doc = fitz.open(output_path)
    remaining_failures = 0
    for pno, page in enumerate(verify_doc):
        fills = _page_fill_rects(page)
        image_rects = _page_image_rects(page)
        for span in _scan_page_spans(page):
            rgb = span["rgb"]
            kind, bg_rgb = _detect_background(span["cx"], span["cy"], fills, image_rects, background)
            if kind == "image":
                continue  # already declared unfixable in the first pass
            ratio = contrast_ratio(rgb, bg_rgb)
            if ratio < effective_target and (pno, rgb) not in declared_unfixable_keys:
                remaining_failures += 1
    verify_doc.close()

    result["ok"] = remaining_failures == 0 and len(result["unfixable"]) == 0
    return result


def main():
    parser = argparse.ArgumentParser(description="WCAG 1.4.3 text-contrast remediation for PDFs (background-aware)")
    parser.add_argument("input", help="Input PDF path")
    parser.add_argument("-o", "--output", help="Output PDF path (default: INPUT.contrast_fixed.pdf)")
    parser.add_argument("--target-ratio", type=float, default=4.5, help="Target contrast ratio (default 4.5)")
    parser.add_argument(
        "--safety-margin", type=float, default=DEFAULT_SAFETY_MARGIN,
        help=f"Effective target floor for headroom against third-party checkers (default {DEFAULT_SAFETY_MARGIN})",
    )
    args = parser.parse_args()

    result = remediate_contrast(args.input, args.output, args.target_ratio, safety_margin=args.safety_margin)
    print(json.dumps(result, indent=2))
    sys.exit(0 if result["ok"] else 1)


if __name__ == "__main__":
    main()
