"""
preprocess.py
Two preprocessing pipelines:
- for_trocr: gentle enhancement only — TrOCR needs natural-looking
  greyscale images, NOT binarized. Binarization breaks it completely.
- for_donut: same gentle pipeline — Donut also expects natural images.
The aggressive binarization from the previous version was causing
TrOCR to output "0 0" instead of text.
"""

import cv2
import numpy as np
from PIL import Image


def _pil_to_cv2(img: Image.Image) -> np.ndarray:
    arr = np.array(img.convert("RGB"))
    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)


def _cv2_to_pil(arr: np.ndarray) -> Image.Image:
    rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)


def deskew(cv_img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
    gray = cv2.bitwise_not(gray)
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    coords = np.column_stack(np.where(thresh > 0))
    if coords.shape[0] < 50:
        return cv_img
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    if abs(angle) < 0.5 or abs(angle) > 20:
        return cv_img
    (h, w) = cv_img.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(
        cv_img, M, (w, h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE
    )


def gentle_enhance(cv_img: np.ndarray) -> np.ndarray:
    """
    Gentle pipeline safe for TrOCR and Donut:
    1. Mild denoise
    2. CLAHE contrast boost on L channel only
    3. NO binarization — keeps natural appearance models expect
    """
    # Mild denoise
    cv_img = cv2.fastNlMeansDenoisingColored(cv_img, None, 3, 3, 7, 21)

    # CLAHE on L channel only
    lab = cv2.cvtColor(cv_img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    cv_img = cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2BGR)

    return cv_img


def upscale_if_small(cv_img: np.ndarray, min_height: int = 1200) -> np.ndarray:
    h, w = cv_img.shape[:2]
    if h < min_height:
        scale = min_height / h
        cv_img = cv2.resize(
            cv_img, (int(w * scale), min_height),
            interpolation=cv2.INTER_CUBIC
        )
    return cv_img


def preprocess_image(image: Image.Image, do_deskew: bool = True) -> Image.Image:
    """Main pipeline — gentle enhance only, NO binarization."""
    cv_img = _pil_to_cv2(image)
    cv_img = upscale_if_small(cv_img)
    cv_img = gentle_enhance(cv_img)
    if do_deskew:
        cv_img = deskew(cv_img)
    return _cv2_to_pil(cv_img)