"""Document detection and 4-point perspective transforms.

Functions for detecting documents in photos (bright paper on dark background)
and applying perspective transforms using detected corner points.
"""

import logging

import cv2
import numpy as np

logger = logging.getLogger(__name__)


def _refine_corners_with_edges(
    gray: np.ndarray,
    corners: np.ndarray,
    doc_mask: np.ndarray,
    contour: np.ndarray | None = None,
) -> np.ndarray:
    """Refine document corners by fitting lines to the brightness-contour sides.

    When brightness-based corner detection clips to image edges (because the
    document extends to or beyond the frame), we can still recover precise
    corners.  The brightness mask contour already follows the document border
    — we just need to compute where its sides actually intersect instead of
    using the approximated polygon vertices.

    Algorithm:
      1. Split the contour into 4 sides using the initial corners as delimiters.
      2. Fit a line to each side using cv2.fitLine (robust L2 fit).
      3. Compute corners as intersections of adjacent side-lines.
      4. Validate the result (convex, reasonable area).

    Args:
        gray:     Grayscale image
        corners:  4×2 initial corners [TL, TR, BR, BL]
        doc_mask: Binary mask of the bright document region
        contour:  Optional full contour from the brightness detection.  If
                  None, the contour is re-extracted from *doc_mask*.

    Returns:
        Refined 4×2 corners, or the original corners if refinement fails.
    """
    h, w = gray.shape[:2]
    edge_margin = 10  # pixels from image edge = "touching boundary"

    # Only refine when at least one corner touches the image boundary
    on_edge = _corners_on_image_edge(corners, w, h, edge_margin)
    if not np.any(on_edge):
        return corners  # All corners are well inside the image

    # ---- 1. Get the contour ----
    pts = _document_contour_points(doc_mask, contour)
    if pts is None:
        return corners

    # ---- 2. Find contour-index of points closest to each corner ----
    sorted_ci, sorted_corner_idx = _sorted_corner_indices(pts, corners)

    # ---- 3. Split contour into 4 sides and fit lines ----
    fitted_lines = _fit_contour_side_lines(_contour_sides(pts, sorted_corner_idx))

    # ---- 4. Compute corners as intersections of adjacent lines ----
    refined, n_refined = _refined_edge_corners(corners, on_edge, sorted_ci, fitted_lines, w, h)
    if n_refined <= 0:
        return corners

    if _valid_refined_quad(refined, w, h):
        logger.info(
            f"Contour-refined {n_refined} corner(s): "
            f"TL=({refined[0, 0]:.0f},{refined[0, 1]:.0f}), "
            f"TR=({refined[1, 0]:.0f},{refined[1, 1]:.0f}), "
            f"BR=({refined[2, 0]:.0f},{refined[2, 1]:.0f}), "
            f"BL=({refined[3, 0]:.0f},{refined[3, 1]:.0f})"
        )
        return refined

    return corners


def _corners_on_image_edge(
    corners: np.ndarray,
    width: int,
    height: int,
    edge_margin: int,
) -> np.ndarray:
    return (
        (corners[:, 0] < edge_margin)
        | (corners[:, 0] > width - edge_margin)
        | (corners[:, 1] < edge_margin)
        | (corners[:, 1] > height - edge_margin)
    )


def _document_contour_points(
    doc_mask: np.ndarray,
    contour: np.ndarray | None,
) -> np.ndarray | None:
    if contour is None:
        contours, _ = cv2.findContours(doc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if not contours:
            return None
        contour = max(contours, key=cv2.contourArea)

    assert contour is not None
    pts = contour.reshape(-1, 2).astype(np.float64)
    return pts if len(pts) >= 20 else None


def _sorted_corner_indices(
    pts: np.ndarray,
    corners: np.ndarray,
) -> tuple[list[int], list[int]]:
    corner_indices = []
    for corner in corners:
        dists = np.linalg.norm(pts - corner.astype(np.float64), axis=1)
        corner_indices.append(int(np.argmin(dists)))

    sorted_ci = sorted(range(4), key=lambda k: corner_indices[k])
    return sorted_ci, [corner_indices[k] for k in sorted_ci]


def _contour_sides(pts: np.ndarray, sorted_corner_idx: list[int]) -> list[np.ndarray]:
    sides: list[np.ndarray] = []
    for i in range(4):
        start = sorted_corner_idx[i]
        end = sorted_corner_idx[(i + 1) % 4]
        sides.append(_contour_side(pts, start, end))
    return sides


def _contour_side(pts: np.ndarray, idx_start: int, idx_end: int) -> np.ndarray:
    if idx_end > idx_start:
        return pts[idx_start : idx_end + 1]
    return np.vstack([pts[idx_start:], pts[: idx_end + 1]])


def _fit_contour_side_lines(
    sides: list[np.ndarray],
) -> list[tuple[float, float, float, float] | None]:
    fitted_lines: list[tuple[float, float, float, float] | None] = []
    for side_pts in sides:
        fitted_lines.append(_fit_contour_side_line(side_pts))
    return fitted_lines


def _fit_contour_side_line(side_pts: np.ndarray) -> tuple[float, float, float, float] | None:
    if len(side_pts) < 5:
        return None
    if len(side_pts) > 200:
        side_pts = side_pts[:: len(side_pts) // 200]
    line = cv2.fitLine(side_pts.astype(np.float32), cv2.DIST_L2, 0, 0.01, 0.01)
    return float(line[0][0]), float(line[1][0]), float(line[2][0]), float(line[3][0])


def _refined_edge_corners(
    corners: np.ndarray,
    on_edge: np.ndarray,
    sorted_ci: list[int],
    fitted_lines: list[tuple[float, float, float, float] | None],
    width: int,
    height: int,
) -> tuple[np.ndarray, int]:
    refined = corners.copy()
    n_refined = 0
    for i in range(4):
        orig_ci = sorted_ci[i]
        if not on_edge[orig_ci]:
            continue
        l_before = fitted_lines[(i - 1) % 4]
        l_after = fitted_lines[i]
        if l_before is None or l_after is None:
            continue
        point = _line_intersect_from_fit(l_before, l_after, width, height)
        if point is not None:
            refined[orig_ci] = np.array(point, dtype=np.float32)
            n_refined += 1
    return refined, n_refined


def _line_intersect_from_fit(
    l1: tuple[float, float, float, float],
    l2: tuple[float, float, float, float],
    width: int,
    height: int,
) -> tuple[float, float] | None:
    vx1, vy1, x01, y01 = l1
    vx2, vy2, x02, y02 = l2
    denom = vx1 * vy2 - vy1 * vx2
    if abs(denom) < 1e-9:
        return None
    dx = x02 - x01
    dy = y02 - y01
    t = (dx * vy2 - dy * vx2) / denom
    ix = x01 + t * vx1
    iy = y01 + t * vy1
    if -width * 0.15 <= ix <= width * 1.15 and -height * 0.15 <= iy <= height * 1.15:
        return ix, iy
    return None


def _valid_refined_quad(refined: np.ndarray, width: int, height: int) -> bool:
    if not _is_convex_quad(refined):
        logger.debug("Refined quad is non-convex — keeping original corners")
        return False
    if cv2.contourArea(refined) <= height * width * 0.15:
        logger.debug("Refined quad too small — keeping original corners")
        return False
    return True


def _is_convex_quad(q: np.ndarray) -> bool:
    """Whether the four corners turn the same way all the way round.

    The 2-D cross product is written out because NumPy 2.0 removed it from
    ``np.cross``, which now raises for two-element vectors rather than
    returning the scalar z-component. Every call raised, so quad refinement was
    never validated and the exception took the whole perspective cascade down
    with it -- silently, since the caller falls back to the original image.
    """
    for i in range(4):
        origin = q[i]
        a = q[(i + 1) % 4] - origin
        b = q[(i + 2) % 4] - origin
        if a[0] * b[1] - a[1] * b[0] <= 0:
            return False
    return True


def detect_photo_document_borders(
    image: np.ndarray, margin_px: int = 20, min_area_ratio: float = 0.15
) -> np.ndarray | None:
    """
    Detect document borders in a photographed document using brightness analysis.

    Detects bright paper on a dark surface by:
    1. Checking if borders are significantly darker than the center (photo indicator)
    2. Thresholding on brightness to isolate the document region
    3. Finding the largest bright contour and approximating it to a quadrilateral
    4. Refining corners with edge detection when they clip to image boundaries
    5. Adding a safety margin around the detected corners

    Args:
        image: Input BGR image
        margin_px: Safety margin (pixels) to add around detected corners
        min_area_ratio: Minimum area of detected region as ratio of image area

    Returns:
        4x2 numpy array of ordered corner points [top-left, top-right,
        bottom-right, bottom-left], or None if document not detected
    """
    h, w = image.shape[:2]
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # ========== STEP 1: Check if this looks like a photo (dark borders, bright center) ==========
    border_w = max(w // 10, 20)
    border_h = max(h // 10, 20)

    border_means = [
        gray[:border_h, :].mean(),  # top
        gray[-border_h:, :].mean(),  # bottom
        gray[:, :border_w].mean(),  # left
        gray[:, -border_w:].mean(),  # right
    ]
    center_brightness = gray[h // 4 : 3 * h // 4, w // 4 : 3 * w // 4].mean()

    # For a true photo (document on desk), at least 3 of 4 borders must be
    # significantly darker than center. Scans of bound documents often have
    # only 1 dark border (the binding shadow), which should NOT trigger this.
    dark_border_threshold = 30
    dark_borders = sum(1 for b in border_means if center_brightness - b > dark_border_threshold)

    if dark_borders < 3:
        logger.debug(
            f"Not a photo: only {dark_borders}/4 borders are dark "
            f"(borders={[f'{b:.0f}' for b in border_means]}, center={center_brightness:.0f})"
        )
        return None

    border_brightness = np.mean(border_means)

    # ========== STEP 2: Find document region using brightness thresholding ==========
    # Threshold halfway between border and center brightness
    thresh_val = int((border_brightness + center_brightness) / 2)
    _, doc_mask = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY)

    # Morphological cleanup to remove noise and fill gaps
    kernel_close = np.ones((15, 15), np.uint8)
    doc_mask = cv2.morphologyEx(doc_mask, cv2.MORPH_CLOSE, kernel_close)
    kernel_open = np.ones((10, 10), np.uint8)
    doc_mask = cv2.morphologyEx(doc_mask, cv2.MORPH_OPEN, kernel_open)

    # Find contours in the brightness mask
    contours, _ = cv2.findContours(doc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        logger.debug("No bright region found in photo")
        return None

    # Get the largest contour (should be the document)
    largest = max(contours, key=cv2.contourArea)
    area = cv2.contourArea(largest)
    image_area = h * w

    if area < image_area * min_area_ratio:
        logger.debug(
            f"Bright region too small: {area / image_area:.1%} of image (need {min_area_ratio:.0%})"
        )
        return None

    # If the bright region fills almost the entire image, this is a scan, not a photo
    if area > image_area * 0.92:
        logger.debug("Bright region fills entire image — this is a scan, not a photo")
        return None

    # ========== STEP 3: Approximate contour to a quadrilateral ==========
    peri = cv2.arcLength(largest, True)
    corners = None

    # Try progressively looser approximation to find a quadrilateral
    for eps_mult in [0.02, 0.04, 0.06, 0.08, 0.10]:
        approx = cv2.approxPolyDP(largest, eps_mult * peri, True)
        if len(approx) == 4:
            corners = order_points(approx.reshape(4, 2).astype(np.float32))
            break

    if corners is None:
        # Fallback: use minimum area rotated rectangle
        rect = cv2.minAreaRect(largest)
        box = cv2.boxPoints(rect)
        corners = order_points(box.astype(np.float32))

    # ========== STEP 4: Refine corners that clip to image edges ==========
    corners = _refine_corners_with_edges(gray, corners, doc_mask)

    # ========== STEP 5: Add safety margin (expand corners outward) ==========
    # Use a proportional margin: 1.5% of the larger dimension, minimum 20px.
    # This ensures enough padding for high-resolution photos where 20px is tiny.
    proportional_margin = max(margin_px, int(max(h, w) * 0.015))
    center = corners.mean(axis=0)
    for i in range(4):
        direction = corners[i] - center
        norm = np.linalg.norm(direction)
        if norm > 0:
            corners[i] += (direction / norm) * proportional_margin

    # NOTE: We intentionally do NOT clamp corners to image bounds.
    # warpPerspective handles out-of-bounds source coordinates by filling
    # with borderValue (white).  Clamping would shift the source quad
    # inward, cropping document content near edges.

    logger.debug(
        f"Photo document detection: "
        f"TL=({corners[0, 0]:.0f}, {corners[0, 1]:.0f}), "
        f"TR=({corners[1, 0]:.0f}, {corners[1, 1]:.0f}), "
        f"BR=({corners[2, 0]:.0f}, {corners[2, 1]:.0f}), "
        f"BL=({corners[3, 0]:.0f}, {corners[3, 1]:.0f}), "
        f"area={area / image_area:.1%}"
    )

    return corners


def correct_photo_perspective(
    image: np.ndarray, corners: np.ndarray, preserve_aspect: bool = True
) -> np.ndarray:
    """
    Apply perspective correction to a photographed document.

    Args:
        image: Input BGR image
        corners: 4x2 array of corners [top-left, top-right, bottom-right, bottom-left]
        preserve_aspect: If True, try to preserve A4/Letter aspect ratio

    Returns:
        Corrected image
    """
    # Calculate dimensions
    top_w = np.linalg.norm(corners[1] - corners[0])
    bot_w = np.linalg.norm(corners[2] - corners[3])
    left_h = np.linalg.norm(corners[3] - corners[0])
    right_h = np.linalg.norm(corners[2] - corners[1])

    new_w = int(max(top_w, bot_w))
    new_h = int(max(left_h, right_h))

    # Adjust to standard paper aspect ratio.
    # Photos taken at steep angles distort the aspect ratio severely, so we
    # use a wide tolerance to catch documents that look nearly square due to
    # extreme perspective foreshortening.
    if preserve_aspect:
        aspect = new_h / new_w if new_w > 0 else 1.0
        A4_ASPECT = 297 / 210  # ~1.414
        LETTER_ASPECT = 11 / 8.5  # ~1.294
        tolerance = 0.5  # wide tolerance for photo distortion

        if abs(aspect - A4_ASPECT) < tolerance:
            new_h = int(new_w * A4_ASPECT)
        elif abs(aspect - LETTER_ASPECT) < tolerance and abs(aspect - LETTER_ASPECT) < abs(
            aspect - A4_ASPECT
        ):
            new_h = int(new_w * LETTER_ASPECT)
        elif abs(aspect - 1 / A4_ASPECT) < tolerance:
            # Landscape A4
            new_w = int(new_h * A4_ASPECT)
        elif abs(aspect - 1 / LETTER_ASPECT) < tolerance:
            # Landscape Letter
            new_w = int(new_h * LETTER_ASPECT)

    # Destination rectangle
    dst = np.array(
        [[0, 0], [new_w - 1, 0], [new_w - 1, new_h - 1], [0, new_h - 1]], dtype=np.float32
    )

    # Perspective transform
    M = cv2.getPerspectiveTransform(corners, dst)
    warped = cv2.warpPerspective(
        image,
        M,
        (new_w, new_h),
        borderValue=(255, 255, 255),  # White background
    )

    logger.info(f"Photo perspective correction applied: {new_w}x{new_h}")
    return warped


def order_points(pts: np.ndarray) -> np.ndarray:
    """
    Order points in clockwise order: top-left, top-right, bottom-right, bottom-left.

    Args:
        pts: Array of 4 points

    Returns:
        Ordered points array
    """
    rect = np.zeros((4, 2), dtype="float32")

    # Top-left has smallest sum, bottom-right has largest sum
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]

    # Top-right has smallest diff, bottom-left has largest diff
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]

    return rect


def four_point_transform(image: np.ndarray, pts: np.ndarray) -> np.ndarray:
    """
    Apply perspective transform to warp the image based on four points.

    Args:
        image: Input image
        pts: Four corner points of the document

    Returns:
        Warped image with corrected perspective
    """
    rect = order_points(pts)
    (tl, tr, br, bl) = rect

    # Compute width (max of top and bottom width)
    width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    max_width = max(int(width_a), int(width_b))

    # Compute height (max of left and right height)
    height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    max_height = max(int(height_a), int(height_b))

    # Destination points
    dst = np.array(
        [[0, 0], [max_width - 1, 0], [max_width - 1, max_height - 1], [0, max_height - 1]],
        dtype="float32",
    )

    # Compute perspective transform matrix and apply
    matrix = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, matrix, (max_width, max_height))

    return warped


def detect_document_contour(image: np.ndarray) -> np.ndarray | None:
    """
    Detect the document contour in an image.

    Uses edge detection and contour finding to locate the document boundaries.

    Args:
        image: Input BGR image

    Returns:
        Four corner points of the document, or None if not found
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Apply Gaussian blur
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # Edge detection
    edges = cv2.Canny(blurred, 50, 150)

    # Dilate to connect edges
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)

    # Find contours
    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    if not contours:
        return None

    # Sort by area, largest first
    contours = sorted(contours, key=cv2.contourArea, reverse=True)

    # Find the largest contour that can be approximated to a quadrilateral
    for contour in contours[:5]:
        # Approximate the contour
        peri = cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, 0.02 * peri, True)

        # If it has 4 points, it's likely the document
        if len(approx) == 4:
            # Check if it's large enough (at least 20% of image area)
            area = cv2.contourArea(approx)
            image_area = image.shape[0] * image.shape[1]
            if area > image_area * 0.2:
                return approx.reshape(4, 2).astype("float32")

    return None


def _contour_needs_perspective_correction(
    contour: np.ndarray,
    threshold: float,
) -> bool:
    """Return whether an already-detected contour exceeds the angular threshold."""

    pts = order_points(contour)
    (tl, tr, br, bl) = pts

    # Calculate angles of each edge
    angles = []

    # Top edge angle
    top_angle = abs(np.arctan2(tr[1] - tl[1], tr[0] - tl[0]))
    angles.append(top_angle)

    # Bottom edge angle
    bottom_angle = abs(np.arctan2(br[1] - bl[1], br[0] - bl[0]))
    angles.append(bottom_angle)

    # Left edge angle (should be ~90 degrees = pi/2)
    left_angle = abs(np.arctan2(bl[1] - tl[1], bl[0] - tl[0]) - np.pi / 2)
    angles.append(left_angle)

    # Right edge angle
    right_angle = abs(np.arctan2(br[1] - tr[1], br[0] - tr[0]) - np.pi / 2)
    angles.append(right_angle)

    max_skew = max(angles)

    logger.debug(f"Perspective analysis: max_skew={max_skew:.4f}, threshold={threshold}")

    return max_skew > threshold
