"""Document/receipt contour detection and perspective correction.

Detection runs two independent candidate generators — Canny edges and Otsu
brightness thresholding — because each fails on different real-world cases
(Canny struggles with low-contrast paper edges; brightness thresholding
struggles when a bright hand or lighting glare merges with the page). Every
candidate contour from both is scored on how "paper-shaped" it is (filled,
convex, plausible aspect ratio) rather than just picking whichever contour
happens to have the largest area, which was previously fooled by fused
hand+background blobs and by fragmented outlines on high-resolution frames.
"""

from __future__ import annotations

from dataclasses import dataclass

import cv2
import numpy as np


@dataclass
class DocumentDetection:
    quad: np.ndarray  # 4x2 float32, corner points in source-image coordinates
    area_ratio: float  # contour area / frame area
    method: str  # "quad" (clean 4-corner polygon) | "min_rect" (fallback)
    source: str  # "edges" | "threshold" — which candidate generator won
    confidence: float  # 0-1 composite shape score (higher = more paper-like)


def _order_points(pts: np.ndarray) -> np.ndarray:
    """Order 4 points as top-left, top-right, bottom-right, bottom-left."""
    pts = pts.astype(np.float32)
    s = pts.sum(axis=1)
    d = np.diff(pts, axis=1).ravel()
    return np.array(
        [
            pts[np.argmin(s)],  # top-left: smallest x+y
            pts[np.argmin(d)],  # top-right: smallest y-x
            pts[np.argmax(s)],  # bottom-right: largest x+y
            pts[np.argmax(d)],  # bottom-left: largest y-x
        ],
        dtype=np.float32,
    )


def _shape_score(contour: np.ndarray, frame_area: float) -> dict | None:
    area = cv2.contourArea(contour)
    area_ratio = area / frame_area
    x, y, w, h = cv2.boundingRect(contour)
    if w == 0 or h == 0:
        return None
    extent = area / (w * h)
    hull = cv2.convexHull(contour)
    hull_area = cv2.contourArea(hull)
    solidity = area / hull_area if hull_area > 0 else 0.0
    aspect = max(w, h) / max(1, min(w, h))
    return {
        "area_ratio": area_ratio,
        "extent": extent,
        "solidity": solidity,
        "aspect": aspect,
    }


def _composite_score(
    stats: dict, min_area_ratio: float, max_area_ratio: float, max_aspect: float
) -> float:
    """Higher is more "paper-like". Negative disqualifies the candidate.

    A page should be a large-ish (but not whole-frame — that's usually
    background/lighting, not a document), filled (high extent), convex
    (high solidity) rectangle. Weighted toward shape over raw size so a
    smaller clean rectangle beats a larger fused hand+background blob.
    """
    if stats["area_ratio"] < min_area_ratio or stats["area_ratio"] > max_area_ratio:
        return -1.0
    if stats["aspect"] > max_aspect:
        return -1.0
    size_term = min(stats["area_ratio"] / 0.5, 1.0)
    return stats["extent"] * 0.45 + stats["solidity"] * 0.35 + size_term * 0.20


def _candidates(mask: np.ndarray, frame_area: float, top_n: int = 8) -> list[np.ndarray]:
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:top_n]
    return contours


def _edge_mask(gray: np.ndarray, canny_low: int, canny_high: int) -> np.ndarray:
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blur, canny_low, canny_high)
    # Kernel scaled to image size so gaps in the outline close reliably
    # regardless of source resolution (a fixed pixel kernel that works at
    # 1000px wide leaves 4x-larger gaps unclosed at 4K).
    k = max(3, round(min(gray.shape) * 0.01)) | 1  # odd
    return cv2.dilate(edges, np.ones((k, k), np.uint8), iterations=2)


def _threshold_mask(gray: np.ndarray) -> np.ndarray:
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    _, mask = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    k = max(3, round(min(gray.shape) * 0.01)) | 1
    kernel = np.ones((k, k), np.uint8)
    return cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)


def _quad_from_contour(contour: np.ndarray) -> tuple[np.ndarray, str]:
    peri = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
    if len(approx) == 4:
        return _order_points(approx.reshape(4, 2)), "quad"
    rect = cv2.minAreaRect(contour)
    return _order_points(cv2.boxPoints(rect)), "min_rect"


def detect_document(
    image_bgr: np.ndarray,
    min_area_ratio: float = 0.04,
    max_area_ratio: float = 0.92,
    max_aspect: float = 6.0,
    canny_low: int = 50,
    canny_high: int = 150,
    detect_width: int = 1000,
) -> DocumentDetection | None:
    """Find the best document-like contour in the frame.

    Contour search runs on a copy downscaled to detect_width (edge/threshold
    gap-closing kernels are sized relative to this, not the source
    resolution), then the winning quad is scaled back up to source
    resolution for a full-quality perspective warp.
    """
    h, w = image_bgr.shape[:2]
    frame_area = float(h * w)

    scale = detect_width / w if w > detect_width else 1.0
    small = (
        cv2.resize(image_bgr, (round(w * scale), round(h * scale)), interpolation=cv2.INTER_AREA)
        if scale != 1.0
        else image_bgr
    )
    small_area = float(small.shape[0] * small.shape[1])
    gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)

    best: tuple[float, np.ndarray, str] | None = None  # (score, contour, source)
    for source, mask in (
        ("edges", _edge_mask(gray, canny_low, canny_high)),
        ("threshold", _threshold_mask(gray)),
    ):
        for contour in _candidates(mask, small_area):
            stats = _shape_score(contour, small_area)
            if stats is None:
                continue
            score = _composite_score(stats, min_area_ratio, max_area_ratio, max_aspect)
            if score < 0:
                continue
            if best is None or score > best[0]:
                best = (score, contour, source)

    if best is None:
        return None

    score, contour, source = best
    small_quad, method = _quad_from_contour(contour)
    quad = scale_quad(small_quad, 1.0 / scale, 1.0 / scale) if scale != 1.0 else small_quad
    area_ratio = cv2.contourArea(contour) / small_area

    return DocumentDetection(
        quad=quad,
        area_ratio=float(area_ratio),
        method=method,
        source=source,
        confidence=float(score),
    )


def warp_document(image_bgr: np.ndarray, quad: np.ndarray) -> np.ndarray:
    """Perspective-correct the quad region into a clean top-down crop."""
    quad = _order_points(np.asarray(quad, dtype=np.float32))
    tl, tr, br, bl = quad

    width = int(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))
    height = int(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))
    width, height = max(width, 1), max(height, 1)

    dst = np.array(
        [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]],
        dtype=np.float32,
    )
    matrix = cv2.getPerspectiveTransform(quad, dst)
    return cv2.warpPerspective(image_bgr, matrix, (width, height))


def scale_quad(quad: np.ndarray, scale_x: float, scale_y: float) -> np.ndarray:
    """Rescale a quad detected on a downscaled frame back to full resolution."""
    quad = np.asarray(quad, dtype=np.float32).copy()
    quad[:, 0] *= scale_x
    quad[:, 1] *= scale_y
    return quad


def pad_quad(
    quad: np.ndarray, image_shape: tuple[int, ...], margin_fraction: float = 0.04
) -> np.ndarray:
    """Expand a quad outward from its centroid by margin_fraction, clamped to the frame.

    approxPolyDP fits a straight-sided polygon to a real paper edge that's
    often slightly curled or wrinkled, so the tightest-fit quad routinely
    clips a sliver of the page (a character or two off the left margin of
    every line, in the worst cases). A small outward pad recovers that
    margin; refine_crop's paper-band tightening trims any extra
    mat/background it pulls in back off afterward.
    """
    h, w = image_shape[:2]
    quad = np.asarray(quad, dtype=np.float32)
    centroid = quad.mean(axis=0)
    padded = centroid + (quad - centroid) * (1.0 + margin_fraction)
    padded[:, 0] = np.clip(padded[:, 0], 0, w - 1)
    padded[:, 1] = np.clip(padded[:, 1], 0, h - 1)
    return padded