from __future__ import annotations

import cv2
import numpy as np


def order_quad_points(pts: np.ndarray) -> np.ndarray:
    pts = np.asarray(pts, dtype=np.float32).reshape(4, 2)
    s = pts.sum(axis=1)
    d = np.diff(pts, axis=1).reshape(-1)

    tl = pts[np.argmin(s)]
    br = pts[np.argmax(s)]
    tr = pts[np.argmin(d)]
    bl = pts[np.argmax(d)]
    return np.array([tl, tr, br, bl], dtype=np.float32)


def detect_page_quad(image_bgr: np.ndarray) -> np.ndarray | None:
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)

    bw = cv2.adaptiveThreshold(
        blur,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        41,
        15,
    )

    contours, _ = cv2.findContours(bw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None

    img_area = image_bgr.shape[0] * image_bgr.shape[1]
    best = None
    best_area = 0.0

    for cnt in contours:
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
        if len(approx) != 4:
            continue

        area = cv2.contourArea(approx)
        if area < img_area * 0.2:
            continue

        if area > best_area:
            best = approx
            best_area = area

    if best is None:
        return None

    return order_quad_points(best.reshape(4, 2))


def warp_quad_to_canvas(
    image_bgr: np.ndarray,
    quad: np.ndarray,
    out_w: int,
    out_h: int,
) -> np.ndarray:
    dst = np.array(
        [[0, 0], [out_w - 1, 0], [out_w - 1, out_h - 1], [0, out_h - 1]],
        dtype=np.float32,
    )
    M = cv2.getPerspectiveTransform(quad.astype(np.float32), dst)
    warped = cv2.warpPerspective(
        image_bgr,
        M,
        (out_w, out_h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE,
    )
    return warped


def preprocess_scanned_image(
    image_bgr: np.ndarray,
    expected_width: int,
    expected_height: int,
) -> tuple[np.ndarray, dict]:
    meta = {
        "used_quad": False,
        "expected_width": int(expected_width),
        "expected_height": int(expected_height),
    }

    quad = detect_page_quad(image_bgr)
    if quad is not None:
        warped = warp_quad_to_canvas(image_bgr, quad, expected_width, expected_height)
        meta["used_quad"] = True
        return warped, meta

    resized = cv2.resize(
        image_bgr,
        (expected_width, expected_height),
        interpolation=cv2.INTER_CUBIC,
    )
    return resized, meta


def rotate_scale_to_canvas(
    image_bgr: np.ndarray,
    angle_deg: float,
    scale: float,
    out_w: int,
    out_h: int,
) -> np.ndarray:
    h, w = image_bgr.shape[:2]
    center = (w / 2.0, h / 2.0)
    M = cv2.getRotationMatrix2D(center, angle_deg, scale)
    warped = cv2.warpAffine(
        image_bgr,
        M,
        (w, h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE,
    )
    if warped.shape[1] != out_w or warped.shape[0] != out_h:
        warped = cv2.resize(warped, (out_w, out_h), interpolation=cv2.INTER_CUBIC)
    return warped


def search_best_candidate(
    base_bgr: np.ndarray,
    expected_width: int,
    expected_height: int,
    candidate_fn,
    angles: list[float] | None = None,
    scales: list[float] | None = None,
) -> dict:
    if angles is None:
        angles = [x / 4.0 for x in range(-8, 9)]  # -2.0 ~ 2.0
    if scales is None:
        scales = [1.0 + x * 0.005 for x in range(-6, 7)]  # 0.97 ~ 1.03

    best = None

    for angle in angles:
        for scale in scales:
            cand = rotate_scale_to_canvas(
                base_bgr,
                angle_deg=angle,
                scale=scale,
                out_w=expected_width,
                out_h=expected_height,
            )
            result = candidate_fn(cand)
            soft = result.get("soft")
            soft_score = float(np.sum(np.abs(soft))) if soft is not None else 0.0
            score = (1e9 if result.get("text_ok") else 0.0) + soft_score

            item = {
                **result,
                "angle": float(angle),
                "scale": float(scale),
                "score": score,
            }

            if best is None or item["score"] > best["score"]:
                best = item

    return best