import io
import cv2
from PIL import Image, ImageOps, ImageFilter
import numpy as np

# Minimum side length below which we upscale the image before sending to OCR.
# Camera photos below this threshold tend to have insufficient pixel density for tiny text.
_MIN_SHORT_SIDE = 1000


class ImageProcessor:
    @staticmethod
    def process_bytes_to_numpy(image_bytes: bytes) -> np.ndarray:
        """
        Converts raw image bytes to an RGB NumPy array ready for PaddleOCR.

        Pipeline:
          1. EXIF auto-rotation  (fixes portrait/landscape mobile photos)
          2. Minimum-resolution guard  (upscale if image is too small)
          3. CLAHE contrast enhancement  (equalises shadows / lighting gradients)
          4. Unsharp masking  (sharpens soft/blurry edges without amplifying noise)
        """
        img = Image.open(io.BytesIO(image_bytes))

        # 1. Auto-rotate based on EXIF orientation tag (critical for mobile photos)
        img = ImageOps.exif_transpose(img)

        # Ensure RGB
        img_rgb = img.convert("RGB")

        # 2. Minimum resolution guard: upscale if the shorter side is too small
        img_rgb = ImageProcessor._ensure_min_resolution(img_rgb)

        img_np = np.array(img_rgb)

        # 3. CLAHE contrast enhancement
        try:
            img_np = ImageProcessor.enhance_contrast(img_np)
        except Exception as e:
            print(f"Warning: Contrast enhancement failed ({e}). Proceeding with raw image.")

        # 4. Unsharp masking for edge sharpening (applied on PIL image, then back to numpy)
        try:
            img_np = ImageProcessor._apply_unsharp_mask(img_np)
        except Exception as e:
            print(f"Warning: Unsharp masking failed ({e}). Skipping sharpening step.")

        return img_np

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _ensure_min_resolution(img: Image.Image) -> Image.Image:
        """
        If the shorter side of the image is below _MIN_SHORT_SIDE pixels,
        scale the image up proportionally using high-quality Lanczos resampling.
        This preserves tiny text that would otherwise be unreadable at low DPI.
        """
        w, h = img.size
        short_side = min(w, h)
        if short_side < _MIN_SHORT_SIDE:
            scale = _MIN_SHORT_SIDE / short_side
            new_w = int(w * scale)
            new_h = int(h * scale)
            img = img.resize((new_w, new_h), Image.LANCZOS)
        return img

    @staticmethod
    def enhance_contrast(img_np: np.ndarray) -> np.ndarray:
        """
        Applies Contrast Limited Adaptive Histogram Equalization (CLAHE)
        to the luminance channel to normalise uneven lighting and shadows.
        """
        yuv = cv2.cvtColor(img_np, cv2.COLOR_RGB2YUV)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        yuv[:, :, 0] = clahe.apply(yuv[:, :, 0])
        return cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB)

    @staticmethod
    def _apply_unsharp_mask(img_np: np.ndarray) -> np.ndarray:
        """
        Applies a gentle unsharp mask to sharpen character edges.

        Parameters tuned conservatively so we enhance detail without
        amplifying noise or creating ringing artefacts on fine print.
        radius=1.5 — small kernel, only sharpens fine detail
        percent=120 — 20% boost in edge contrast
        threshold=3 — only sharpen pixels that differ by at least 3 levels
                       (prevents noise from being sharpened)
        """
        # Work in PIL for built-in UnsharpMask
        pil_img = Image.fromarray(img_np)
        sharpened = pil_img.filter(ImageFilter.UnsharpMask(radius=1.5, percent=120, threshold=3))
        return np.array(sharpened)


image_processor = ImageProcessor()