"""
OCR Processor Module
Applies image preprocessing to improve Tesseract OCR accuracy.
"""
import logging
import re

import cv2
import numpy as np
import pytesseract
from PIL import Image

logger = logging.getLogger(__name__)


class OCRProcessor:
    """
    Wraps Tesseract OCR with image preprocessing pipeline:
    - Grayscale conversion
    - Adaptive thresholding (handles uneven lighting)
    - Noise removal via median blur
    - Deskewing for rotated scans
    """

    def __init__(self, tesseract_cmd: str = None):
        if tesseract_cmd:
            pytesseract.pytesseract.tesseract_cmd = tesseract_cmd

        # Tesseract config: PSM 3 = fully automatic page segmentation
        self.tsv_config = r"--oem 3 --psm 3"

    def preprocess(self, image: Image.Image) -> np.ndarray:
        """
        Apply a preprocessing pipeline to improve OCR quality.
        Returns a numpy array (grayscale, thresholded image).
        """
        # Convert PIL to OpenCV format
        img_array = np.array(image.convert("RGB"))
        img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

        # Step 1: Convert to grayscale
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)

        # Step 2: Denoise
        denoised = cv2.medianBlur(gray, 3)

        # Step 3: Adaptive thresholding (works well for uneven lighting)
        thresh = cv2.adaptiveThreshold(
            denoised,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            31,  # block size
            10,  # constant subtracted from mean
        )

        # Step 4: Morphological dilation to connect broken characters
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))
        processed = cv2.dilate(thresh, kernel, iterations=1)

        return processed

    def extract(self, image: Image.Image) -> str:
        """
        Run OCR on a PIL Image.
        Returns cleaned, normalized text.
        """
        try:
            processed = self.preprocess(image)
            # Convert back to PIL for pytesseract
            processed_pil = Image.fromarray(processed)
            raw_text = pytesseract.image_to_string(processed_pil, config=self.tsv_config)
            return self._clean_ocr_output(raw_text)
        except Exception as e:
            logger.error(f"OCR extraction failed: {e}")
            # Fallback: try without preprocessing
            try:
                raw_text = pytesseract.image_to_string(image, config=self.tsv_config)
                return self._clean_ocr_output(raw_text)
            except Exception as e2:
                logger.error(f"OCR fallback also failed: {e2}")
                return ""

    def _clean_ocr_output(self, text: str) -> str:
        """
        Clean raw OCR output:
        - Remove control characters
        - Normalize whitespace
        - Remove page numbers and artifacts
        - Merge broken hyphenated lines
        """
        if not text:
            return ""

        # Normalize unicode
        text = text.encode("utf-8", errors="ignore").decode("utf-8")

        # Merge hyphenated line breaks (e.g., "docu-\nment" → "document")
        text = re.sub(r"-\n(\w)", r"\1", text)

        # Replace multiple spaces/tabs with single space
        text = re.sub(r"[ \t]+", " ", text)

        # Keep at most 2 consecutive newlines
        text = re.sub(r"\n{3,}", "\n\n", text)

        # Remove lines that are pure noise (single chars, page numbers)
        lines = text.splitlines()
        cleaned_lines = []
        for line in lines:
            stripped = line.strip()
            # Skip very short lines that look like artifacts
            if len(stripped) <= 1 and not stripped.isalpha():
                continue
            cleaned_lines.append(stripped)

        return "\n".join(cleaned_lines).strip()