"""rasterize job (M2): original PDF/images -> optimized page JPEGs + AI copies + blank check.
Payload: {script_id, batch_id}. Idempotent: if pages exist, only re-enqueues missing extractions."""
from __future__ import annotations
import hashlib
import io

import numpy as np
import pymupdf
from PIL import Image

from .. import config, jobs
from ..ids import new_id
from ..records import raise_flag
from ..settings import settings
from ..storage import get_storage


def ink_stats(img: Image.Image) -> tuple[float, int]:
    """(ink ratio, stroke pixels). Ink = pixels clearly darker than this page's own paper tone,
    so faint pencil counts but light show-through does not. Printed ruled lines are removed, and only
    connected strokes of >= 30 px are counted, so dust specks don't make a blank page 'written'."""
    from scipy import ndimage
    g = np.asarray(img.convert("L"), dtype=np.int16)
    h, w = g.shape
    by, bx = int(h * 0.04), int(w * 0.04)
    g = g[by:h - by, bx:w - bx]
    paper = int(np.percentile(g, 90))            # paper tone of this scan
    dark = g < paper - 35
    dark[dark.mean(axis=1) > 0.6, :] = False     # ruled lines (rows)
    dark[:, dark.mean(axis=0) > 0.6] = False     # margin lines (columns)
    labels, n = ndimage.label(dark)
    if n == 0:
        return 0.0, 0
    sizes = np.bincount(labels.ravel())[1:]
    stroke_px = int(sizes[sizes >= 30].sum())
    return float(dark.mean()), stroke_px


def is_blank(img: Image.Image, min_stroke_px: int) -> tuple[bool, float]:
    ratio, stroke_px = ink_stats(img)
    return stroke_px < min_stroke_px, ratio


def encode_jpeg(img: Image.Image, quality: int, target_kb: int) -> bytes:
    """JPEG at `quality`; step quality down (never below 55) until under 1.5 × target."""
    q = quality
    while True:
        buf = io.BytesIO()
        img.save(buf, "JPEG", quality=q, optimize=True, progressive=True)
        data = buf.getvalue()
        if len(data) <= target_kb * 1024 * 1.5 or q <= 55:
            return data
        q -= 10


def _pages_from_original(data: bytes, dpi: int):
    """Yield PIL images, one per page, from a PDF or a single image file."""
    if data[:4] == b"%PDF":
        doc = pymupdf.open(stream=data, filetype="pdf")
        for page in doc:
            pix = page.get_pixmap(dpi=dpi, colorspace=pymupdf.csGRAY)
            yield Image.frombytes("L", (pix.width, pix.height), pix.samples)
    else:
        img = Image.open(io.BytesIO(data))
        img.load()
        yield img.convert("L")


def handle(conn, job) -> dict:
    p = job["payload_json"]
    tenant_id, script_id = job["tenant_id"], p["script_id"]
    script = conn.execute("SELECT * FROM scripts WHERE id=%s AND tenant_id=%s",
                          (script_id, tenant_id)).fetchone()
    if not script:
        raise jobs.FatalError("script not found")
    batch_id = script["batch_id"]
    storage = get_storage()
    existing = conn.execute("SELECT count(*) AS n FROM pages WHERE script_id=%s AND tenant_id=%s",
                            (script_id, tenant_id)).fetchone()["n"]
    if existing == 0:
        original = storage.get(script["original_file_key"])
        base = script["original_file_key"].rsplit("/", 1)[0] + "/pages"
        min_stroke = int(config.get(conn, "image.blank_min_stroke_px", tenant_id) or settings.blank_min_stroke_px)
        n = 0
        for i, img in enumerate(_pages_from_original(original, settings.raster_dpi), start=1):
            n = i
            jpeg = encode_jpeg(img, settings.jpeg_quality, settings.jpeg_target_kb)
            ai = img.copy()
            ai.thumbnail((settings.ai_image_long_edge, settings.ai_image_long_edge))
            ai_jpeg = encode_jpeg(ai, 80, 400)
            key, ai_key = f"{base}/{i:03d}.jpg", f"{base}/{i:03d}.ai.jpg"
            storage.put(key, jpeg, "image/jpeg")
            storage.put(ai_key, ai_jpeg, "image/jpeg")
            blank, ratio = is_blank(img, min_stroke)
            conn.execute(
                """INSERT INTO pages (id, tenant_id, script_id, sort, kind, image_key, ai_image_key, width, height,
                                      dpi, checksum, ink_ratio, status)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                (new_id(), tenant_id, script_id, i, "blank" if blank else None, key, ai_key, img.width,
                 img.height, settings.raster_dpi, hashlib.sha256(jpeg).hexdigest(), round(ratio, 6),
                 "skipped" if blank else "pending"))
        if n == 0:
            raise jobs.FatalError("document has no pages")
        if n > settings.max_pages_per_script:
            raise_flag(conn, tenant_id=tenant_id, batch_id=batch_id, type_="too_many_pages",
                       script_id=script_id, payload={"pages": n, "max": settings.max_pages_per_script})
        conn.execute("UPDATE scripts SET page_count=%s, status='extracting', updated_at=now() WHERE id=%s AND tenant_id=%s",
                     (n, script_id, tenant_id))
        conn.execute("UPDATE batches SET page_count = page_count + %s WHERE id=%s AND tenant_id=%s",
                     (n, batch_id, tenant_id))

    pending = conn.execute("SELECT id FROM pages WHERE script_id=%s AND tenant_id=%s AND status='pending' ORDER BY sort",
                           (script_id, tenant_id)).fetchall()
    for row in pending:
        jobs.enqueue(conn, "extract_page", {"page_id": row["id"], "script_id": script_id, "batch_id": batch_id},
                     tenant_id=tenant_id, priority=job["priority"], dedupe_key=f"extract:{row['id']}",
                     correlation_id=job["correlation_id"])
    if not pending:  # every page blank: go straight to assembly
        jobs.enqueue(conn, "assemble", {"script_id": script_id, "batch_id": batch_id}, tenant_id=tenant_id,
                     priority=job["priority"], dedupe_key=f"assemble:{script_id}")
    return {"pages": existing or n, "extract_jobs": len(pending)}