"""Runnable check for optional PDF shrinking, including what it refuses to do."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pymupdf

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from noteman_slicer.pdfopt import MIN_DPI, optimise  # noqa: E402

A4_PT = (595, 842)


def _pdf(path: Path, width: int, height: int, *, colour: bool = False, bilevel: bool = False):
    """One full-page image of ruled lines, at the given pixel size."""
    art = np.full((height, width, 3), 255, np.uint8)
    for i in range(6):
        y = int(height * (0.2 + i * 0.03))
        art[y : y + max(1, height // 900), int(width * 0.1) : int(width * 0.9)] = 0
    if colour:
        art[: height // 3, :, 0] = 40  # a strong blue cast over the top third
        art[: height // 3, :, 1] = 90
    grey = art[:, :, 0] if not colour else None

    doc = pymupdf.open()
    page = doc.new_page(width=A4_PT[0], height=A4_PT[1])
    if bilevel:
        pix = pymupdf.Pixmap(pymupdf.csGRAY, width, height, bytearray(grey.tobytes()), False)
        page.insert_image(page.rect, pixmap=pix)
        doc.save(path, garbage=4, deflate=True)
        # Re-save through a 1-bit PNG so the stored image really is bilevel.
        import cv2

        ok, buf = cv2.imencode(".png", (grey > 127).astype(np.uint8) * 255)
        doc2 = pymupdf.open()
        p2 = doc2.new_page(width=A4_PT[0], height=A4_PT[1])
        p2.insert_image(p2.rect, stream=buf.tobytes())
        doc2.save(path, garbage=4, deflate=True)
        return
    stream = art if colour else np.dstack([grey] * 3)
    import cv2

    ok, buf = cv2.imencode(".jpg", stream, [cv2.IMWRITE_JPEG_QUALITY, 92])
    page.insert_image(page.rect, stream=buf.tobytes())
    doc.save(path, garbage=4, deflate=True)


def main() -> int:
    tmp = Path(__file__).with_name("_tmp")
    tmp.mkdir(exist_ok=True)

    # A4 is 8.26in wide, so 2480px is ~300 DPI and 800px is ~97 DPI.
    fine, coarse, colour = tmp / "fine.pdf", tmp / "coarse.pdf", tmp / "colour.pdf"
    _pdf(fine, 2480, 3508)
    _pdf(coarse, 800, 1130)
    _pdf(colour, 2480, 3508, colour=True)

    data, report = optimise(pymupdf.open(fine), fine.stat().st_size)
    assert report.converted == 1, report.summary()
    assert data, "a greyscale scan at 300 DPI should shrink"
    assert report.ratio < 0.9, report.ratio
    # The result must still be a readable PDF of the same page count.
    assert len(pymupdf.open(stream=data, filetype="pdf")) == 1

    # Too coarse: staff lines would break, so it is left alone.
    _, report = optimise(pymupdf.open(coarse), coarse.stat().st_size)
    assert report.converted == 0, report.summary()
    assert any("DPI" in reason for reason in report.skipped), report.skipped

    # Genuine colour: artwork is not thrown away.
    _, report = optimise(pymupdf.open(colour), colour.stat().st_size)
    assert report.converted == 0, report.summary()
    assert any("colour" in reason for reason in report.skipped), report.skipped

    # A no-op run reports honestly rather than returning something bigger.
    empty = pymupdf.open()
    empty.new_page()
    data, report = optimise(empty, 1)
    assert data == b"" and report.converted == 0
    assert report.ratio == 1.0

    assert MIN_DPI >= 150, "the floor exists to protect thin staff lines"

    for f in (fine, coarse, colour):
        f.unlink(missing_ok=True)
    tmp.rmdir()
    print("ok")
    return 0


if __name__ == "__main__":
    sys.exit(main())