"""Standard coarse-to-fine splitting: cheap positive content gate, precise fold.

The whole-page gate only establishes obvious writing. All uncertain/sparse pages
use the reference blank detector. Full-resolution fold processing is restricted
to a central strip; edge-reaching overflow and ambiguous ownership fall back to
the complete reference analysis. Full-spread analysis remains an internal fallback.
"""

from __future__ import annotations
import cv2
import numpy as np
import pymupdf
from copienator.auto_page_split import analyse_ink, _analyse_page_full, plan_masks


def quick_page_content(rgb, dpi, target_dpi=75, *, native=True, _fallback=None):
    """Positive-only coarse gate. Ambiguous/sparse pages use full detection."""
    h, w = rgb.shape[:2]
    small = cv2.resize(
        rgb,
        (round(w * target_dpi / dpi), round(h * target_dpi / dpi)),
        interpolation=cv2.INTER_AREA,
    )
    px = target_dpi / 25.4
    lo = small.min(axis=2)
    ch = small.max(axis=2) - lo
    gray = cv2.cvtColor(small, cv2.COLOR_RGB2GRAY)
    strong = (ch > 100) | (gray < 95)
    n, lab, stats, _ = cv2.connectedComponentsWithStats(strong.astype("uint8"), 8)
    ok = np.zeros(n, bool)
    x, y, bw, bh, area = stats[1:].T
    # Only compact interior pen components can establish an obvious written page.
    ok[1:] = (
        (area >= 0.25 * px * px)
        & (bw >= 0.35 * px)
        & (bh >= 0.5 * px)
        & (bw < 20 * px)
        & (bh < 15 * px)
        & (x > 0.09 * small.shape[1])
        & (x + bw < 0.91 * small.shape[1])
    )
    selected = ok[lab]
    grouped = cv2.dilate(
        selected.astype("uint8"),
        np.ones((max(1, round(px)), max(1, round(3 * px))), np.uint8),
    )
    ng, groups, gs, _ = cv2.connectedComponentsWithStats(grouped, 8)
    comp_group = np.zeros(n, np.int32)
    ys, xs = np.nonzero(selected)
    comp_group[lab[ys, xs]] = groups[ys, xs]
    counts = np.bincount(comp_group[ok], minlength=ng)
    amounts = np.bincount(groups[selected], minlength=ng)
    good = (
        (counts >= 4)
        & (amounts >= 3 * px * px)
        & (gs[:, 2] >= 8 * px)
        & (gs[:, 3] >= 2 * px)
    )
    good[0] = False
    if good.any():
        # Retain all strong ink for ownership support; the gate alone is interior-only.
        return (
            strong
            if native
            else cv2.resize(
                strong.astype("uint8"), (w, h), interpolation=cv2.INTER_NEAREST
            )
            > 0
        ), True
    return (_fallback() if _fallback is not None else analyse_ink(rgb, dpi)[0]), False


def _a4_plan(mask, gate, h, w, dpi, gate_dpi):
    count = int(mask.sum())
    return dict(
        format="a4",
        size_px=[w, h],
        outputs=[
            dict(
                side="whole",
                rect_px=[0, 0, w, h],
                state="content" if count else "blank",
                keep=bool(count),
                ink_pixels=count,
                content_dpi=gate_dpi if gate else dpi,
            )
        ],
        crossings=[],
        review_reasons=[],
        extensions_mm=[0, 0],
        analysis=dict(
            mode="fast",
            early_written=[gate],
            fold_fraction=0,
            gate_dpi=gate_dpi,
            fallback=None,
        ),
    )


def analyse_page_fast(rgb, dpi, kind, *, fold_fraction=0.1, gate_dpi=75):
    """Fast analysis for an existing raster image."""
    if kind not in {"a3", "a4"}:
        return _analyse_page_full(rgb, dpi, kind)
    gate_dpi = min(gate_dpi, dpi)
    h, w = rgb.shape[:2]
    if kind == "a4":
        mask, gate = quick_page_content(rgb, dpi, gate_dpi, native=True)
        return _a4_plan(mask, gate, h, w, dpi, gate_dpi)
    middle = w // 2
    halves = []
    early = []
    for a, b in ((0, middle), (middle, w)):
        mask, gate = quick_page_content(rgb[:, a:b], dpi, gate_dpi, native=True)
        halves.append(mask)
        early.append(gate)
    return _assemble_fast(
        halves,
        early,
        h,
        w,
        dpi,
        fold_fraction,
        gate_dpi,
        lambda x0, x1: rgb[:, x0:x1],
        lambda: _analyse_page_full(rgb, dpi, kind),
    )


def _assemble_fast(
    halves, early, h, w, dpi, fold_fraction, gate_dpi, render_roi, reference
):
    middle = w // 2
    half = max(1, round(w * fold_fraction / 2))
    x0 = max(0, middle - half)
    x1 = min(w, middle + half)
    roi = render_roi(x0, x1)
    content = np.zeros(roi.shape[:2], bool)
    for i, (a, b) in enumerate(((0, middle), (middle, w))):
        lo = max(a, x0)
        hi = min(b, x1)
        mask = halves[i]
        ys = np.arange(h) * mask.shape[0] // h
        xs = (np.arange(lo, hi) - a) * mask.shape[1] // (b - a)
        content[:, lo - x0 : hi - x0] = mask[ys[:, None], xs[None, :]]
    ink, _ = analyse_ink(roi, dpi, spread=True)
    remaining = {}
    plan = plan_masks(
        ink, dpi, "a3", content_ink=content, rgb=roi, _remaining_content=remaining
    )
    # If a proposed overflow reaches the strip's edge or ownership remains
    # ambiguous, use the complete reference analysis rather than truncate it.
    margin = round(2 * dpi / 25.4)
    if (
        plan["review_reasons"]
        or plan["outputs"][0]["rect_px"][2] >= x1 - x0 - margin
        or plan["outputs"][1]["rect_px"][0] <= margin
    ):
        plan = reference()
        plan["analysis"] = dict(mode="fast", fallback="full-spread")
        return plan
    for i, (a, b) in enumerate(((0, middle), (middle, w))):
        mask = halves[i]
        side = "left" if i == 0 else "right"
        if not early[i]:
            lo = max(a, x0)
            hi = min(b, x1)
            mask[:, lo - a : hi - a] = remaining[side]
            count = int(mask.sum())
        else:
            # Exact count of the implied nearest-neighbour full-resolution mask,
            # without allocating it. Coarse masks establish obvious content only.
            wy = np.diff(
                np.ceil(np.arange(mask.shape[0] + 1) * h / mask.shape[0]).astype(
                    "int64"
                )
            )
            wx = np.diff(
                np.ceil(np.arange(mask.shape[1] + 1) * (b - a) / mask.shape[1]).astype(
                    "int64"
                )
            )
            count = int(np.sum(mask * wy[:, None] * wx[None, :]))
        o = plan["outputs"][i]
        o["ink_pixels"] = count
        o["keep"] = count > 0
        o["state"] = "content" if count else "blank"
        o["rect_px"] = (
            [0, 0, o["rect_px"][2] + x0, h]
            if i == 0
            else [o["rect_px"][0] + x0, 0, w, h]
        )
    for c in plan["crossings"]:
        c["box"][0] += x0
        c["box"][2] += x0
        for box in c.get("anchor_boxes", []):
            if box is not None:
                box[0] += x0
                box[2] += x0
    plan["size_px"] = [w, h]
    plan["analysis"] = dict(
        mode="fast",
        gate_dpi=gate_dpi,
        fold_fraction=fold_fraction,
        early_written=early,
        fallback=None,
    )
    return plan


def analyse_pdf_page_fast(page, dpi, kind, *, fold_fraction=0.1, gate_dpi=75):
    """Render a small overview and the precise fold directly from the PDF.

    Full-resolution halves are rendered only when the positive overview gate
    cannot establish writing. Raster clips are aligned to the same integer
    pixel grid as a complete reference render, including rotated CropBoxes.
    """
    gate_dpi = min(gate_dpi, dpi)
    bounds = (page.rect * pymupdf.Matrix(dpi / 72, dpi / 72)).irect
    h, w = bounds.height, bounds.width
    coarse_pix = page.get_pixmap(
        dpi=round(gate_dpi), colorspace=pymupdf.csRGB, alpha=False
    )
    coarse = np.frombuffer(coarse_pix.samples, np.uint8).reshape(
        coarse_pix.height, coarse_pix.width, 3
    )

    def render(x0, x1):
        rect = pymupdf.Rect(x0 * 72 / dpi, 0, x1 * 72 / dpi, page.rect.height)
        pix = page.get_pixmap(
            dpi=round(dpi), clip=rect, colorspace=pymupdf.csRGB, alpha=False
        )
        data = np.frombuffer(pix.samples, np.uint8).reshape(pix.height, pix.width, 3)
        out = np.full((h, x1 - x0, 3), 255, np.uint8)
        left = max(x0, pix.x)
        top = max(0, pix.y)
        right = min(x1, pix.x + pix.width)
        bottom = min(h, pix.y + pix.height)
        out[top:bottom, left - x0 : right - x0] = data[
            top - pix.y : bottom - pix.y, left - pix.x : right - pix.x
        ]
        return out

    def reference():
        return _analyse_page_full(render(0, w), dpi, kind)

    if kind not in {"a3", "a4"}:
        return reference(), coarse
    if kind == "a4":
        mask, gate = quick_page_content(
            coarse,
            gate_dpi,
            gate_dpi,
            native=True,
            _fallback=lambda: analyse_ink(render(0, w), dpi)[0],
        )
        return _a4_plan(mask, gate, h, w, dpi, gate_dpi), coarse
    middle = w // 2
    cm = coarse.shape[1] // 2
    halves = []
    early = []
    for a, b, ca, cb in ((0, middle, 0, cm), (middle, w, cm, coarse.shape[1])):
        mask, gate = quick_page_content(
            coarse[:, ca:cb],
            gate_dpi,
            gate_dpi,
            native=True,
            _fallback=lambda a=a, b=b: analyse_ink(render(a, b), dpi)[0],
        )
        halves.append(mask)
        early.append(gate)
    plan = _assemble_fast(
        halves, early, h, w, dpi, fold_fraction, gate_dpi, render, reference
    )
    plan["analysis"]["rendering"] = "direct-pdf-roi"
    return plan, coarse