"""Repro: get_page_image's 1.5x supersample-downsample loses borderline OCR text.

Self-contained: generates a degraded scanned-page PDF (image-only), then
compares RapidOCR text detection on (a) a direct pdfium render at scale 3 vs
(b) docling's PyPdfiumDocumentBackend.get_page_image(scale=3) of the same page.

Requires: docling, rapidocr, onnxruntime, pypdfium2, pillow, numpy.
"""

import io
from pathlib import Path

import numpy as np
import pypdfium2 as pdfium
from PIL import Image, ImageFilter

OUT = Path("/tmp/supersample-repro.pdf")
PAGE_W, PAGE_H = 612, 792


def make_source_pdf() -> bytes:
    lines = [
        (24, 710, "Accessible Documents Report"),
        (12, 660, "This is the first paragraph of the sample document. It exists to give"),
        (12, 644, "the layout analysis model a block of running body text to classify,"),
        (12, 628, "spanning several lines so it reads as a paragraph rather than a label."),
        (12, 584, "A second paragraph follows after a vertical gap. It is visually distinct"),
        (12, 568, "from the first and should be detected as a separate text element."),
        (12, 524, "- tagging quality"),
        (12, 508, "- reading order"),
        (12, 492, "- alt text coverage"),
    ]
    content = b"".join(
        f"BT /F1 {s} Tf 72 {y} Td ({t}) Tj ET\n".encode() for s, y, t in lines
    )
    objects = [
        b"<< /Type /Catalog /Pages 2 0 R >>",
        b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
        f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {PAGE_W} {PAGE_H}] "
        f"/Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>".encode(),
        b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
        f"<< /Length {len(content)} >>\nstream\n".encode() + content + b"\nendstream",
    ]
    return assemble(objects)


def assemble(objects: list[bytes]) -> bytes:
    out = bytearray(b"%PDF-1.4\n")
    offsets = []
    for num, body in enumerate(objects, start=1):
        offsets.append(len(out))
        out += f"{num} 0 obj\n".encode() + body + b"\nendobj\n"
    xref = len(out)
    out += f"xref\n0 {len(objects) + 1}\n".encode() + b"0000000000 65535 f \n"
    for off in offsets:
        out += f"{off:010d} 00000 n \n".encode()
    out += f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n".encode()
    return bytes(out)


def make_scanned_pdf() -> None:
    """Render the source, degrade like a 95-dpi noisy scan, embed as image-only PDF."""
    src = Path("/tmp/supersample-src.pdf")
    src.write_bytes(make_source_pdf())
    pdf = pdfium.PdfDocument(src)
    clean = pdf[0].render(scale=300 / 72.0).to_pil().convert("L")
    pdf.close()
    small = clean.resize(
        (int(clean.width * 95 / 300), int(clean.height * 95 / 300)),
        Image.Resampling.BILINEAR,
    ).filter(ImageFilter.GaussianBlur(radius=0.6))
    rng = np.random.default_rng(1001)
    px = np.clip(np.asarray(small, dtype=np.float32) + rng.normal(0, 16.0, (small.height, small.width)), 0, 255)
    noisy = Image.fromarray(px.astype(np.uint8), mode="L")
    buf = io.BytesIO()
    noisy.save(buf, format="JPEG", quality=35)
    jpeg = buf.getvalue()
    content = f"q {PAGE_W} 0 0 {PAGE_H} 0 0 cm /Im0 Do Q".encode()
    objects = [
        b"<< /Type /Catalog /Pages 2 0 R >>",
        b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
        f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {PAGE_W} {PAGE_H}] "
        f"/Resources << /XObject << /Im0 4 0 R >> >> /Contents 5 0 R >>".encode(),
        f"<< /Type /XObject /Subtype /Image /Width {noisy.width} /Height {noisy.height} "
        f"/ColorSpace /DeviceGray /BitsPerComponent 8 /Filter /DCTDecode "
        f"/Length {len(jpeg)} >>\nstream\n".encode() + jpeg + b"\nendstream",
        f"<< /Length {len(content)} >>\nstream\n".encode() + content + b"\nendstream",
    ]
    OUT.write_bytes(assemble(objects))


def detect_lines(image: Image.Image) -> list[str]:
    from rapidocr import RapidOCR

    result = RapidOCR()(np.array(image.convert("RGB")))
    return list(result.txts) if result is not None and result.txts is not None else []


def main() -> None:
    make_scanned_pdf()

    # (a) direct pdfium render at the target scale
    pdf = pdfium.PdfDocument(OUT)
    direct = pdf[0].render(scale=3).to_pil()
    pdf.close()

    # (b) docling's backend render of the same page at the same scale
    from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.document import InputDocument

    in_doc = InputDocument(
        path_or_stream=OUT, format=InputFormat.PDF, backend=PyPdfiumDocumentBackend
    )
    backend = in_doc._backend
    docling_img = backend.load_page(0).get_page_image(scale=3)

    a = detect_lines(direct)
    b = detect_lines(docling_img)
    print(f"direct pdfium render @3x:        {len(a)} text lines detected")
    print(f"docling get_page_image(scale=3): {len(b)} text lines detected")
    needle = "second"  # the paragraph line that goes missing
    print(f"line 'A second paragraph...' -> direct: "
          f"{any(needle in t.lower() for t in a)}, docling render: "
          f"{any(needle in t.lower() for t in b)}")


if __name__ == "__main__":
    main()