from aix import jobs
from aix.main import run_one
from aix.pipeline import rasterize

from .conftest import add_script
from .helpers import make_pdf


def test_rasterize_pages_blank_detection_and_fanout(conn, storage, world):
    sid = add_script(conn, storage, world, make_pdf(["text", "blank", "ruled", "faint", "showthrough", "shortfaint"]))
    jobs.enqueue(conn, "rasterize", {"script_id": sid, "batch_id": world["batch"]}, tenant_id=world["tenant"])
    assert run_one(conn, {"rasterize": rasterize.handle})
    pages = conn.execute("SELECT sort, kind, status, image_key, ai_image_key, ink_ratio FROM pages WHERE script_id=%s ORDER BY sort",
                         (sid,)).fetchall()
    kinds = {p["sort"]: p["kind"] for p in pages}
    assert kinds[2] == "blank" and kinds[3] == "blank" and kinds[5] == "blank"   # blank, ruled-only, show-through
    assert kinds[1] is None and kinds[4] is None and kinds[6] is None           # text, faint, short faint kept
    for p in pages:
        size_kb = len(storage.get(p["image_key"])) / 1024
        assert size_kb < 450, (p["sort"], size_kb)
        assert storage.exists(p["ai_image_key"])
    ex = conn.execute("SELECT count(*) AS n FROM jobs WHERE type='extract_page' AND status='queued'").fetchone()["n"]
    assert ex == 3
    assert conn.execute("SELECT page_count FROM scripts WHERE id=%s", (sid,)).fetchone()["page_count"] == 6


def test_rasterize_is_idempotent(conn, storage, world):
    sid = add_script(conn, storage, world, make_pdf(["text"]))
    job = {"payload_json": {"script_id": sid, "batch_id": world["batch"]}, "tenant_id": world["tenant"],
           "priority": 100, "correlation_id": "c"}
    with conn.transaction():
        rasterize.handle(conn, job)
    with conn.transaction():
        rasterize.handle(conn, job)
    assert conn.execute("SELECT count(*) AS n FROM pages WHERE script_id=%s", (sid,)).fetchone()["n"] == 1
    # dedupe key prevents a second live extract job for the same page
    assert conn.execute("SELECT count(*) AS n FROM jobs WHERE type='extract_page'").fetchone()["n"] == 1


def test_dust_specks_on_grey_paper_are_blank():
    import numpy as np
    from PIL import Image
    rng = np.random.default_rng(1)
    a = np.clip(228 + rng.integers(-6, 7, (2339, 1654)), 0, 255).astype(np.uint8)
    for _ in range(300):
        y, x = rng.integers(100, 2200), rng.integers(100, 1550)
        a[y:y + 3, x:x + 3] = 90
    blank, _ = rasterize.is_blank(Image.fromarray(a), 120)
    assert blank