"""Group per-frame document detections into capture events.

Stable-window-only detection misses documents shown while hands are still
adjusting them — motion scores stay above threshold even though a receipt
is clearly visible and detectable frame-by-frame. This clusters consecutive
frames where a document contour was found (with a time-gap to split separate
placements) and keeps the best frame from each cluster.

"Best" means hands-free first, sharp second: fingers holding a page flat are
usually present in most frames of an event but absent in at least a few
(right after placing / before flipping). A sharpness-only pick routinely
selects a hand-in-frame frame, which then has to be inpainted (lossy).
Preferring the lowest-skin frame keeps the original pixels instead.

Memory: scoring runs as frames arrive; only tiny gray proxies + metrics are
retained for clustering. Full-res BGR is dropped immediately so a long 4K
video cannot OOM by holding thousands of detections in RAM (callers re-decode
winning frame indices in a third pass).
"""

from __future__ import annotations

from dataclasses import dataclass

import cv2
import numpy as np

from paperpod.capture.video import Frame
from paperpod.vision.document import DocumentDetection
from paperpod.vision.motion import motion_score, prepare_gray
from paperpod.vision.refine import skin_mask
from paperpod.vision.sharpness import sharpness_score

# Downscale width for per-frame skin analysis; full-res 4K would dominate
# pass-2 runtime for no accuracy gain.
_SKIN_ANALYSIS_WIDTH = 400

# Skin coverage below this fraction of the document region counts as a
# "clean" (hands-free) frame; between the two thresholds counts as "touched"
# (fingertip on an edge); above is "held" (hand across the page).
_SKIN_CLEAN = 0.003
_SKIN_TOUCHED = 0.03

# Reject a "cleaner" frame if it's this dull relative to the sharpest
# sample seen in the same event.
_SHARP_FLOOR = 0.4

# Long placements that include a page flip keep the same paper outline, so
# contour clustering never splits them. Instead we look for *quiet dwells*
# (runs of near-zero inter-sample motion) inside a long event — each dwell
# is one page lying flat. Page 2 of the Jan '25 TD statement sat still for
# ~0.7s around t=10.4; hand-wobble stretches between dwells never qualify.
_QUIET_MOTION = 0.015
_MIN_DWELL_S = 0.55
_LONG_EVENT_S = 4.0
_MOTION_SIZE = (160, 90)


@dataclass
class DocumentEvent:
    event_id: int
    t_start: float
    t_end: float
    sample_count: int
    best_frame_t: float
    best_frame_index: int
    sharpness: float
    skin_fraction: float
    detection: DocumentDetection
    # Populated by the pipeline after a re-decode pass; None right after
    # clustering so the scorer never retains full-res images.
    frame: Frame | None = None


@dataclass
class ScoredSample:
    """Per-detection metrics without the full-res BGR frame."""

    index: int
    t: float
    shape: tuple[int, ...]  # (h, w, c) of the source frame
    det: DocumentDetection
    sharp: float
    skin: float
    gray_tiny: np.ndarray  # for page-flip motion between samples


def _quad_region(image: np.ndarray, quad: np.ndarray) -> np.ndarray:
    """Bounding-box crop of the detected document within the frame."""
    h, w = image.shape[:2]
    x0 = max(0, int(quad[:, 0].min()))
    x1 = min(w, int(np.ceil(quad[:, 0].max())))
    y0 = max(0, int(quad[:, 1].min()))
    y1 = min(h, int(np.ceil(quad[:, 1].max())))
    if x1 - x0 >= 8 and y1 - y0 >= 8:
        return image[y0:y1, x0:x1]
    return image


def skin_fraction_in_quad(image: np.ndarray, quad: np.ndarray) -> float:
    """Fraction of the detected document region that looks like skin.

    Measured on the quad's bounding box (downscaled), so a hand resting on
    the mat next to the document doesn't count — only skin over the paper.
    """
    region = _quad_region(image, quad)
    rh, rw = region.shape[:2]
    if rw > _SKIN_ANALYSIS_WIDTH:
        region = cv2.resize(
            region, (_SKIN_ANALYSIS_WIDTH, max(1, round(rh * _SKIN_ANALYSIS_WIDTH / rw)))
        )
    return float((skin_mask(region) > 0).mean())


def document_sharpness(image: np.ndarray, quad: np.ndarray) -> float:
    """Sharpness of the document region only.

    Whole-frame Laplacian variance is dominated by high-frequency background
    (wood grain), so a blurry mid-flip frame that exposes more table can
    outscore a crisp flat page. Scoring inside the detection quad measures
    the thing we actually keep.
    """
    return sharpness_score(_quad_region(image, quad))


def score_sample(frame: Frame, det: DocumentDetection) -> ScoredSample:
    """Extract clustering metrics and drop the full-res pixels."""
    small = cv2.resize(frame.image, _MOTION_SIZE)
    return ScoredSample(
        index=frame.index,
        t=frame.t,
        shape=frame.image.shape,
        det=det,
        sharp=document_sharpness(frame.image, det.quad),
        skin=skin_fraction_in_quad(frame.image, det.quad),
        gray_tiny=prepare_gray(small, blur_ksize=5),
    )


def _skin_bucket(fraction: float) -> int:
    if fraction < _SKIN_CLEAN:
        return 0
    if fraction < _SKIN_TOUCHED:
        return 1
    return 2


def _same_document(
    prev: DocumentDetection,
    curr: DocumentDetection,
    frame_shape: tuple[int, ...],
    area_tolerance: float = 0.12,
    centroid_tolerance: float = 0.12,
) -> bool:
    """True if two detections likely belong to the same physical document.

    On a black mat, the empty table is often detected as a large contour on
    almost every frame, which bridges separate receipts/forms into one giant
    cluster if we only use time gaps. A sudden change in area or position
    means a new document was placed (Costco receipt area ~0.26 vs mat ~0.65).
    """
    if abs(prev.area_ratio - curr.area_ratio) > area_tolerance:
        return False
    h, w = frame_shape[:2]
    prev_c = prev.quad.mean(axis=0)
    curr_c = curr.quad.mean(axis=0)
    if float(np.linalg.norm(prev_c - curr_c)) > centroid_tolerance * max(w, h):
        return False
    return True


def _plausible_capture(det: DocumentDetection) -> bool:
    """Skip mat/table false positives that sit between receipt and full-page sizes.

    On a black mat, empty-table contours usually land around area_ratio
    0.45-0.57. Real receipts cluster below 0.40; full tax forms above 0.58.
    """
    ar = det.area_ratio
    return ar <= 0.40 or ar >= 0.58


def _split_on_page_flips(
    samples: list[ScoredSample], min_samples: int
) -> list[list[ScoredSample]]:
    """Pull quiet dwells out of a long continuous contour (page flips).

    Returns one segment per flat-page dwell when two or more qualify;
    otherwise returns the original samples unchanged so ordinary
    single-page placements aren't shredded by brief hand motion.
    """
    if len(samples) < 4:
        return [samples]
    duration = samples[-1].t - samples[0].t
    if duration < _LONG_EVENT_S:
        return [samples]

    motions = [0.0]
    for i in range(1, len(samples)):
        motions.append(
            motion_score(samples[i - 1].gray_tiny, samples[i].gray_tiny, pixel_threshold=12)
        )

    dwells: list[list[ScoredSample]] = []
    start: int | None = None
    for i, sample in enumerate(samples):
        quiet = motions[i] <= _QUIET_MOTION
        if quiet and start is None:
            start = i
        elif not quiet and start is not None:
            dwell = samples[start:i]
            if dwell[-1].t - dwell[0].t >= _MIN_DWELL_S and len(dwell) >= min_samples:
                dwells.append(dwell)
            start = None
    if start is not None:
        dwell = samples[start:]
        if dwell[-1].t - dwell[0].t >= _MIN_DWELL_S and len(dwell) >= min_samples:
            dwells.append(dwell)

    # Only commit to multi-page split when we actually found 2+ quiet dwells.
    return dwells if len(dwells) >= 2 else [samples]


class _Cluster:
    """Accumulates samples for one placement; picks the best frame on flush.

    Preference among frames ≥ 40% as sharp as the cluster's sharpest sample:
    lowest skin bucket, then highest sharpness. Soft clean frames never
    outrank sharp lightly-touched ones.

    First/last samples of long events are ignored — they are usually the
    hand placing or lifting the page (high Laplacian streaks, low ink),
    which is how page 7 of a statement once lost to its own lift frame.
    """

    def __init__(self, sample: ScoredSample) -> None:
        self.t_start = sample.t
        self.t_end = sample.t
        self.samples: list[ScoredSample] = [sample]

    @property
    def count(self) -> int:
        return len(self.samples)

    def add(self, sample: ScoredSample) -> None:
        self.t_end = sample.t
        self.samples.append(sample)

    def pick_best(self, samples: list[ScoredSample] | None = None) -> ScoredSample:
        samples = samples if samples is not None else self.samples
        pick_from = samples
        if len(pick_from) >= 5:
            # Drop the placing/lifting edges.
            pick_from = pick_from[1:-1]
        max_sharp = max(s.sharp for s in pick_from)
        floor = max_sharp * _SHARP_FLOOR
        eligible = [s for s in pick_from if s.sharp >= floor] or pick_from
        return min(eligible, key=lambda s: (_skin_bucket(s.skin), -s.sharp))


def find_document_events(
    samples: list[tuple[Frame, DocumentDetection]] | list[ScoredSample],
    gap_s: float = 2.0,
    min_samples: int = 2,
) -> list[DocumentEvent]:
    """Cluster consecutive detections; split when gap_s passes with no detection.

    Long clusters (page kept under the camera while flipped) are further
    split on motion peaks into separate dwells so each page gets its own
    capture. samples must be in time order.

    Accepts either (Frame, detection) tuples (scores immediately, drops BGR)
    or pre-built ScoredSample list from the streaming pipeline.
    """
    if not samples:
        return []

    scored: list[ScoredSample]
    if isinstance(samples[0], ScoredSample):
        scored = list(samples)  # type: ignore[arg-type]
    else:
        scored = [score_sample(frame, det) for frame, det in samples]  # type: ignore[misc]

    clusters: list[_Cluster] = []
    active: _Cluster | None = None

    def flush() -> None:
        nonlocal active
        if active is not None and active.count >= min_samples:
            clusters.append(active)
        active = None

    for sample in scored:
        if active is None:
            active = _Cluster(sample)
            continue
        if sample.t - active.t_end > gap_s or not _same_document(
            active.samples[-1].det, sample.det, sample.shape
        ):
            flush()
            active = _Cluster(sample)
            continue
        active.add(sample)

    flush()

    events: list[DocumentEvent] = []
    event_id = 0
    for cluster in clusters:
        for segment in _split_on_page_flips(cluster.samples, min_samples):
            best = cluster.pick_best(segment)
            events.append(
                DocumentEvent(
                    event_id=event_id,
                    t_start=segment[0].t,
                    t_end=segment[-1].t,
                    sample_count=len(segment),
                    best_frame_t=best.t,
                    best_frame_index=best.index,
                    sharpness=best.sharp,
                    skin_fraction=round(best.skin, 4),
                    detection=best.det,
                    frame=None,
                )
            )
            event_id += 1
    return events