"""Runnable check for source classification and raster loading.

Builds its own PDFs so it needs no corpus files (scores are copyrighted and
gitignored). Run with `python tests/test_pdf.py`.
"""

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.pdf import SourceType, open_source, page_raster  # noqa: E402

A4 = pymupdf.paper_rect("a4")


def _vector_pdf(path: Path, pages: int = 2) -> None:
    doc = pymupdf.open()
    for _ in range(pages):
        page = doc.new_page(width=A4.width, height=A4.height)
        page.draw_line((50, 100), (A4.width - 50, 100))
        page.insert_text((50, 150), "notation", fontsize=24)
    doc.save(path)


def _scan_pdf(path: Path, pages: int = 2, w: int = 1653, h: int = 2332, rotation: int = 0) -> None:
    """Each page is one full-page grayscale image — what a real scan looks like.

    `rotation` reproduces a sheet fed sideways: the image is stored in its own
    orientation and /Rotate turns it upright for a viewer.
    """
    art = np.full((h, w), 255, np.uint8)
    art[500:505, 100 : w - 100] = 0  # a staff line, so it isn't uniform
    art[:60, :60] = 0  # a corner mark, so orientation is checkable
    pix = pymupdf.Pixmap(pymupdf.csGRAY, w, h, bytearray(art.tobytes()), False)
    doc = pymupdf.open()
    for _ in range(pages):
        page = doc.new_page(width=A4.width, height=A4.width * h / w)
        page.insert_image(page.rect, pixmap=pix)
        page.set_rotation(rotation)
    doc.save(path)


def main() -> int:
    tmp = Path(__file__).with_name("_tmp")
    tmp.mkdir(exist_ok=True)
    vec, scan = tmp / "vector.pdf", tmp / "scan.pdf"
    _vector_pdf(vec)
    _scan_pdf(scan)

    src = open_source(vec)
    assert src.type is SourceType.VECTOR, src.type
    assert not src.overridden
    page = page_raster(src, 0)
    # Rendered at 600 DPI, so an A4 page is ~4960px wide.
    assert page.ndim == 2 and page.dtype == np.uint8, (page.ndim, page.dtype)
    assert 4900 < page.shape[1] < 5000, page.shape
    src.close()

    src = open_source(scan)
    assert src.type is SourceType.RASTER, src.type
    page = page_raster(src, 0)
    # Native resolution of the embedded image, NOT a 600 DPI re-render.
    assert page.shape == (2332, 1653), page.shape
    assert page.min() == 0 and page.max() == 255, (page.min(), page.max())
    src.close()

    # The corner mark sits top-left in an upright scan.
    assert page[:60, :60].max() == 0 and page[:60, -60:].min() == 255

    # A sideways scan comes back upright: the page's /Rotate applies to the
    # image extracted by xref, which bypasses it. Okular gets this right and
    # the slicer used to not.
    sideways = tmp / "sideways.pdf"
    _scan_pdf(sideways, pages=1, w=2332, h=1653, rotation=90)
    src = open_source(sideways)
    assert src.type is SourceType.RASTER
    turned = page_raster(src, 0)
    assert turned.shape == (2332, 1653), turned.shape
    # Turned clockwise, so the mark that was top-left is now top-right.
    assert turned[:60, -60:].max() == 0 and turned[:60, :60].min() == 255
    src.close()

    # An override must win over detection, and say so.
    src = open_source(scan, SourceType.VECTOR)
    assert src.type is SourceType.VECTOR and src.detected is SourceType.RASTER
    assert src.overridden
    assert page_raster(src, 0).shape[1] > 4000, "override must force a render"
    src.close()

    for f in (vec, scan, sideways):
        f.unlink()
    tmp.rmdir()
    print("ok")
    return 0


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