"""Optional shrinking of the original PDF carried in a bundle.

Scanned scores are usually black ink on white paper stored as 8-bit greyscale
or RGB, which costs several times what the same page costs as a bilevel image.
Converting them is worth 7–9× on a real corpus.

Two things it must not do, both found by looking at output rather than at
numbers:

  * A page that is genuinely coloured — cover artwork — loses its artwork.
  * A scan too coarse to have more than about one pixel per staff line comes
    back with the staff lines broken.

Both are detectable before converting, so both are skipped. Everything skipped
is reported, so a caller can say what was left alone and why.

This affects only the archival copy of the score. Slices are cut from the
original before any of this and are unchanged either way.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import cv2
import numpy as np
import pymupdf

# Below this many pixels per inch as the image is *placed on the page*, staff
# lines are about a pixel wide and thresholding breaks them. Measured against a
# corpus where the one failure sat at ~115 DPI and the successes at 260+.
MIN_DPI = 200

# An image is "coloured" when this share of sampled pixels are off-grey by
# more than _CHROMA. The two populations are far apart: measured on a corpus,
# cover artwork sits at 44% while a greyscale scan's sensor tint reaches 3%.
# Ten percent sits in the gap with room on both sides.
_CHROMA = 24
_COLOUR_SHARE = 0.10

_BLOCK = 31  # adaptive threshold window
_OFFSET = 15


@dataclass
class Report:
    before: int = 0
    after: int = 0
    converted: int = 0
    skipped: dict[str, int] = field(default_factory=dict)

    @property
    def ratio(self) -> float:
        return self.after / self.before if self.before else 1.0

    def skip(self, reason: str) -> None:
        self.skipped[reason] = self.skipped.get(reason, 0) + 1

    def summary(self) -> str:
        if not self.converted:
            return "nothing to optimise — every image is already bilevel, coloured or too coarse"
        parts = [
            f"{self.before / 1024:.0f} KB → {self.after / 1024:.0f} KB "
            f"({self.ratio * 100:.0f}%), {self.converted} images converted"
        ]
        for reason, count in sorted(self.skipped.items()):
            parts.append(f"{count} left alone: {reason}")
        return "\n".join(parts)


def _is_coloured(image: np.ndarray) -> bool:
    if image.ndim != 3 or image.shape[2] < 3:
        return False
    sample = image[::4, ::4, :3].astype(np.int16)
    spread = sample.max(axis=2) - sample.min(axis=2)
    return float((spread > _CHROMA).mean()) > _COLOUR_SHARE


def _placed_dpi(page: pymupdf.Page, item, width: int) -> float:
    """Pixels per inch of an image as it appears on the page.

    Not the pixel count: a page split into tiles has small images at a high
    resolution, and a full-page image can be large yet coarse.
    """
    try:
        bbox = pymupdf.Rect(page.get_image_bbox(item))
    except (ValueError, RuntimeError):
        return float("inf")
    inches = abs(bbox.width) / 72.0
    return width / inches if inches > 0 else float("inf")


def optimise(doc: pymupdf.Document, source_bytes: int) -> tuple[bytes, Report]:
    """Return the optimised PDF and a report of what was done.

    `doc` is modified in place, so pass a copy or reopen afterwards.
    """
    report = Report(before=source_bytes)

    for page in doc:
        for item in page.get_images(full=True):
            xref = item[0]
            info = doc.extract_image(xref)

            if info.get("bpc") == 1:
                report.skip("already bilevel")
                continue

            if _placed_dpi(page, item, info["width"]) < MIN_DPI:
                report.skip(f"below {MIN_DPI} DPI, staff lines would break")
                continue

            raw = cv2.imdecode(np.frombuffer(info["image"], np.uint8), cv2.IMREAD_UNCHANGED)
            if raw is None:
                report.skip("unreadable encoding")
                continue

            if _is_coloured(raw):
                report.skip("coloured artwork")
                continue

            gray = cv2.cvtColor(raw, cv2.COLOR_BGR2GRAY) if raw.ndim == 3 else raw
            bilevel = cv2.adaptiveThreshold(
                gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, _BLOCK, _OFFSET
            )
            ok, buffer = cv2.imencode(".png", bilevel, [cv2.IMWRITE_PNG_COMPRESSION, 9])
            if not ok:
                report.skip("re-encoding failed")
                continue
            try:
                page.replace_image(xref, stream=buffer.tobytes())
            except (ValueError, RuntimeError):
                report.skip("could not be replaced")
                continue
            report.converted += 1

    data = doc.tobytes(garbage=4, deflate=True, clean=True)
    # Never hand back something larger than what came in.
    if len(data) >= source_bytes:
        report.after = source_bytes
        report.converted = 0
        report.skip("no saving available")
        return b"", report

    report.after = len(data)
    return data, report


def preview(original: pymupdf.Document, optimised: pymupdf.Document, dpi: int = 260):
    """A stacked before/after crop of the first page, for eyeballing the result.

    The numbers cannot show the failure this guards against — a broken staff
    line is obvious at a glance and invisible in a byte count.
    """
    rect = original[0].rect
    clip = pymupdf.Rect(
        rect.x0 + rect.width * 0.08,
        rect.y0 + rect.height * 0.20,
        rect.x0 + rect.width * 0.58,
        rect.y0 + rect.height * 0.33,
    )

    def render(doc: pymupdf.Document) -> np.ndarray:
        pix = doc[0].get_pixmap(dpi=dpi, clip=clip)
        image = np.frombuffer(pix.samples, np.uint8).reshape(pix.height, pix.width, pix.n)
        return image[:, :, :3] if pix.n >= 3 else cv2.cvtColor(image[:, :, 0], cv2.COLOR_GRAY2BGR)

    before, after = render(original), render(optimised)
    h = min(before.shape[0], after.shape[0])
    w = min(before.shape[1], after.shape[1])
    divider = np.full((4, w, 3), 128, np.uint8)
    return np.vstack([before[:h, :w], divider, after[:h, :w]])