"""Auto-rotate a document crop to right-side-up before OCR/PDF export.

Handles the "I placed it upside down" case that broke naming on the real T4/
T4A test recordings: images looked fine, but OCR read garbage because the
form was rotated 180 degrees.
"""

from __future__ import annotations

from dataclasses import dataclass

import cv2
import numpy as np
import pytesseract

_ROTATIONS = {
    0: None,
    90: cv2.ROTATE_90_CLOCKWISE,
    180: cv2.ROTATE_180,
    270: cv2.ROTATE_90_COUNTERCLOCKWISE,
}


@dataclass
class OrientationResult:
    rotation: int  # degrees clockwise applied: 0, 90, 180, or 270
    method: str  # "osd" | "confidence_sweep" | "none"


def apply_rotation(image_bgr: np.ndarray, rotation: int) -> np.ndarray:
    code = _ROTATIONS.get(rotation % 360)
    return image_bgr if code is None else cv2.rotate(image_bgr, code)


def _downscale(image_bgr: np.ndarray, max_dim: int = 1000) -> np.ndarray:
    """Shrink before feeding Tesseract: orientation only needs word shapes,
    not per-character resolution, and _confidence_sweep_rotation runs this
    up to 4x per document — at full 4K-crop resolution that made rotation
    detection by far the slowest part of the whole pipeline."""
    h, w = image_bgr.shape[:2]
    longest = max(h, w)
    if longest <= max_dim:
        return image_bgr
    scale = max_dim / longest
    return cv2.resize(image_bgr, (round(w * scale), round(h * scale)), interpolation=cv2.INTER_AREA)


def _osd_rotation(image_bgr: np.ndarray, lang: str) -> int | None:
    """Tesseract's orientation-and-script-detection: fast and purpose-built.

    Raises on pages with too little recognizable text (e.g. a mostly-blank
    or badly-cropped receipt sliver); callers should fall back in that case.
    Note: OSD's own confidence scores run on a small, model-dependent scale
    (values well under 10 are common even when correct) so we trust any
    result it returns rather than gating on a confidence threshold.
    """
    try:
        osd = pytesseract.image_to_osd(
            _downscale(image_bgr), lang=lang, output_type=pytesseract.Output.DICT
        )
    except pytesseract.TesseractError:
        return None
    return int(osd.get("rotate", 0)) % 360


def _confidence_sweep_rotation(
    image_bgr: np.ndarray, lang: str, psm: int, min_words: int = 3, min_word_len: int = 3
) -> int:
    """Fallback: OCR all 4 rotations, keep the one with the most confident text.

    Used when OSD can't commit (too little text). Scores by *total*
    confidence summed over recognized words at least min_word_len long.

    Both pieces matter and were tuned against real failures, not guessed:
      - Summing (not averaging) confidence: on noisy real-world photos
        (worn thermal print, low-contrast background), mean confidence
        stays within a similar narrow band at every rotation, but the
        correct orientation recognizes many more words than a sideways one.
        Mean-confidence-only picked a 90-degrees-off rotation (9 words) over
        the correct one (23 words) on a real low-contrast receipt.
      - Filtering to words >= min_word_len: sideways/upside-down *clean*
        synthetic text tends to fragment into dozens of single-character
        glyph artifacts that Tesseract reports with deceptively high
        individual confidence, which otherwise wins the sum outright.
        Requiring 3+ characters drops almost all of that noise while barely
        touching real multi-character words.
    """
    small = _downscale(image_bgr)
    best_rotation = 0
    best_score = -1.0
    for rotation in (0, 90, 180, 270):
        candidate = apply_rotation(small, rotation)
        data = pytesseract.image_to_data(
            candidate, lang=lang, config=f"--psm {psm}", output_type=pytesseract.Output.DICT
        )
        confidences = [
            float(c)
            for c, t in zip(data["conf"], data["text"])
            if c not in ("-1", -1) and len(t.strip()) >= min_word_len
        ]
        if len(confidences) < min_words:
            continue
        total_conf = sum(confidences)
        if total_conf > best_score:
            best_score = total_conf
            best_rotation = rotation
    return best_rotation


def detect_orientation(image_bgr: np.ndarray, lang: str = "eng", psm: int = 6) -> OrientationResult:
    """Figure out how many degrees clockwise to rotate image_bgr to be upright."""
    rotation = _osd_rotation(image_bgr, lang)
    if rotation is not None:
        return OrientationResult(rotation=rotation, method="osd")
    rotation = _confidence_sweep_rotation(image_bgr, lang, psm)
    return OrientationResult(rotation=rotation, method="confidence_sweep")


def auto_rotate(
    image_bgr: np.ndarray, lang: str = "eng", psm: int = 6
) -> tuple[np.ndarray, OrientationResult]:
    """Detect and correct orientation. Returns (rotated_image, result)."""
    result = detect_orientation(image_bgr, lang=lang, psm=psm)
    return apply_rotation(image_bgr, result.rotation), result