"""PDF input: classify a score source and hand back page rasters.

Two source types, never mixed within one PDF (docs/spec.md):

  raster  — a scan; every page carries one full-page image, and *that image
            is the scan*. It is extracted at its native resolution rather
            than re-rendered: real scans in this corpus run ~200 DPI, and
            re-rendering at 600 would triple the pixel count for no detail.
  vector  — an engraving; nothing to extract, so the page is rendered.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import numpy as np
import pymupdf

VECTOR_RENDER_DPI = 600

# An image covering at least this fraction of the page is the page's scan
# rather than an illustration sitting on an engraving.
_FULL_PAGE_AREA = 0.5


class SourceType(Enum):
    RASTER = "raster"
    VECTOR = "vector"


@dataclass
class Source:
    path: Path
    doc: pymupdf.Document
    type: SourceType
    detected: SourceType
    render_dpi: int = VECTOR_RENDER_DPI

    @property
    def overridden(self) -> bool:
        """True when the user's choice disagrees with detection."""
        return self.type is not self.detected

    def __len__(self) -> int:
        return len(self.doc)

    def close(self) -> None:
        self.doc.close()


def _full_page_image(page: pymupdf.Page) -> int | None:
    """xref of the image covering this page, or None."""
    page_area = abs(page.rect.get_area())
    if page_area <= 0:
        return None
    # full=True is required, or get_image_bbox rejects the item.
    for item in page.get_images(full=True):
        try:
            bbox = pymupdf.Rect(page.get_image_bbox(item))
        except ValueError:
            continue
        if abs(bbox.get_area()) >= page_area * _FULL_PAGE_AREA:
            return item[0]
    return None


def classify(doc: pymupdf.Document) -> SourceType:
    """Detection only — the caller confirms with the user (ADR 0004)."""
    scanned = sum(_full_page_image(page) is not None for page in doc)
    return SourceType.RASTER if scanned * 2 > len(doc) else SourceType.VECTOR


def open_source(path: str | Path, source_type: SourceType | None = None) -> Source:
    """Open a PDF. `source_type` overrides detection; it never silently wins."""
    path = Path(path)
    doc = pymupdf.open(path)
    detected = classify(doc)
    return Source(path=path, doc=doc, type=source_type or detected, detected=detected)


def page_raster(source: Source, index: int) -> np.ndarray:
    """One page as a grayscale array, at the resolution the pipeline should work at."""
    page = source.doc[index]

    if source.type is SourceType.RASTER:
        xref = _full_page_image(page)
        if xref is not None:
            # Pixmap(doc, xref) rather than decoding extract_image() bytes:
            # MuPDF handles JBIG2 and CCITT, which no image library will.
            pix = pymupdf.Pixmap(source.doc, xref)
            return _to_gray(pix)
        # A scanned PDF whose page has no embedded image (a blank, or a
        # cover typeset in vector). Rendering is the only option left.

    return _to_gray(page.get_pixmap(dpi=source.render_dpi, colorspace=pymupdf.csGRAY))


def _to_gray(pix: pymupdf.Pixmap) -> np.ndarray:
    if pix.alpha or pix.colorspace is None or pix.colorspace.n != 1:
        pix = pymupdf.Pixmap(pymupdf.csGRAY, pix)
    return np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width)