services/ocr_engine.py

"""
OCR Engine — Extracts text from images and scanned documents using Tesseract.

Uses an adaptive preprocessing strategy:
  1. Try Tesseract with gentle preprocessing (grayscale + light cleanup)
  2. If confidence is poor, retry with medium preprocessing (+ Otsu binarization)
  3. Keep the best result

Tesseract is used for ALL OCR (both images and scanned PDFs).
EasyOCR has been removed to keep memory usage within HuggingFace free-tier limits.
"""
import cv2
import numpy as np
from pathlib import Path
from PIL import Image
import pytesseract
from services.image_preprocessor import (
    preprocess_gentle,
    preprocess_medium,
    preprocess_heavy,
    preprocess_image,
    fix_orientation,
)

# --- Add Windows Tesseract fallback for local testing ---
import os
import sys
if sys.platform == 'win32':
    try:
        pytesseract.get_tesseract_version()
    except Exception:
        for path in [
            r"C:\Program Files\Tesseract-OCR\tesseract.exe",
            r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
            r"C:\Users\kmthu\AppData\Local\Programs\Tesseract-OCR\tesseract.exe",
        ]:
            if os.path.exists(path):
                pytesseract.pytesseract.tesseract_cmd = path
                break
# --------------------------------------------------------


# Minimum average confidence (0-1 scale) to accept a Tesseract result
MIN_ACCEPTABLE_CONFIDENCE = 0.40


def _tesseract_ocr_page(cleaned_img: np.ndarray) -> dict:
    """
    Run Tesseract on a single preprocessed image and return structured results.

    Returns dict with 'text', 'blocks', 'avg_confidence', 'word_count'.
    """
    # Get raw text
    raw_text = pytesseract.image_to_string(cleaned_img, config='--psm 3').strip()

    # Get word-level data for confidence and bounding boxes
    data = pytesseract.image_to_data(
        cleaned_img, output_type=pytesseract.Output.DICT, config='--psm 3'
    )

    blocks = []
    total_confidence = 0.0
    word_count = 0

    for i in range(len(data['text'])):
        text = data['text'][i].strip()
        conf = data['conf'][i]

        try:
            conf_val = float(conf) / 100.0
        except (ValueError, TypeError):
            conf_val = -1.0

        if text and conf_val >= 0:
            x = data['left'][i]
            y = data['top'][i]
            w = data['width'][i]
            h = data['height'][i]
            blocks.append({
                "text": text,
                "confidence": round(conf_val, 3),
                "bbox": [[x, y], [x + w, y], [x + w, y + h], [x, y + h]],
            })
            total_confidence += conf_val
            word_count += 1

    avg_confidence = total_confidence / word_count if word_count > 0 else 0.0

    return {
        "text": raw_text,
        "blocks": blocks,
        "avg_confidence": round(avg_confidence, 3),
        "word_count": word_count,
    }


def _is_gibberish(text: str, avg_confidence: float) -> bool:
    """
    Heuristic check: is the OCR output likely gibberish?
    """
    if avg_confidence < MIN_ACCEPTABLE_CONFIDENCE:
        return True

    if not text or len(text) < 10:
        return True

    # Check ratio of letter/digit characters vs total
    alpha_count = sum(1 for c in text if c.isalnum())
    total_chars = len(text.replace(' ', '').replace('\n', ''))
    if total_chars > 0:
        alpha_ratio = alpha_count / total_chars
        if alpha_ratio < 0.5:
            return True

    return False


def _adaptive_ocr(tmp_path: str, pil_img: Image.Image = None) -> dict:
    """
    Run Tesseract with adaptive preprocessing on a single page/image.

    Tries gentle first, then medium if result is gibberish.
    Returns the best result dict with 'text', 'blocks', 'avg_confidence'.
    """
    best_result = None

    # Load and fix orientation ONCE
    if pil_img is not None:
        # Convert PIL to OpenCV BGR format
        open_cv_image = np.array(pil_img)
        # Handle different modes (RGB, RGBA, L)
        if len(open_cv_image.shape) == 3 and open_cv_image.shape[2] == 3:
            img_bgr = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
        elif len(open_cv_image.shape) == 3 and open_cv_image.shape[2] == 4:
            img_bgr = cv2.cvtColor(open_cv_image, cv2.COLOR_RGBA2BGR)
        else:
            img_bgr = cv2.cvtColor(open_cv_image, cv2.COLOR_GRAY2BGR)
    else:
        img_bgr = cv2.imread(tmp_path)
        if img_bgr is None:
            return {"text": "", "blocks": [], "avg_confidence": 0.0, "word_count": 0}

    oriented_img = fix_orientation(img_bgr)

    # Tier 1: Gentle preprocessing (no binarization — lets Tesseract decide)
    try:
        gentle_img = preprocess_gentle(oriented_img)
        result = _tesseract_ocr_page(gentle_img)
        best_result = result
    except Exception as e:
        print(f"  Gentle preprocessing failed: {e}")

    # Tier 2: Medium preprocessing (+ Otsu binarization) — only if gentle was bad
    if best_result is None or _is_gibberish(best_result['text'], best_result['avg_confidence']):
        try:
            medium_img = preprocess_medium(oriented_img)
            result = _tesseract_ocr_page(medium_img)
            if best_result is None or result['avg_confidence'] > best_result['avg_confidence']:
                best_result = result
        except Exception as e:
            print(f"  Medium preprocessing failed: {e}")

    # Tier 3: Raw grayscale (no preprocessing) — absolute fallback
    if best_result is None or _is_gibberish(best_result['text'], best_result['avg_confidence']):
        try:
            raw_img = cv2.cvtColor(oriented_img, cv2.COLOR_BGR2GRAY)
            if raw_img is not None:
                # Upscale if small
                h, w = raw_img.shape[:2]
                if w < 2000:
                    scale = 2000 / w
                    raw_img = cv2.resize(
                        raw_img, (int(w * scale), int(h * scale)),
                        interpolation=cv2.INTER_CUBIC
                    )
                result = _tesseract_ocr_page(raw_img)
                if best_result is None or result['avg_confidence'] > best_result['avg_confidence']:
                    best_result = result
        except Exception as e:
            print(f"  Raw OCR failed: {e}")

    if best_result is None:
        best_result = {"text": "", "blocks": [], "avg_confidence": 0.0, "word_count": 0}

    return best_result


def extract_from_image(image_path: str | Path, preprocess: bool = True) -> dict:
    """
    Extract text from a standalone image file (JPG, PNG) using Tesseract.
    """
    image_path = str(image_path)

    result = _adaptive_ocr(image_path)

    return {
        "raw_text": result["text"],
        "blocks": result["blocks"],
        "confidence_avg": result["avg_confidence"],
        "page_count": 1,
    }


def extract_from_scanned_pdf(pdf_path: str | Path) -> dict:
    """
    Extract text from a scanned PDF using Tesseract with adaptive preprocessing.

    For each page: converts to image, tries gentle then medium preprocessing,
    keeps the best result.
    """
    try:
        import pdfplumber
        import tempfile
        import os

        all_text_parts = []
        all_blocks = []
        total_confidence = 0.0
        total_detections = 0
        page_count = 0

        with pdfplumber.open(str(pdf_path)) as pdf:
            page_count = len(pdf.pages)

            for page in pdf.pages:
                # Convert PDF page to image at high resolution
                page_image = page.to_image(resolution=300)
                pil_img = page_image.original

                # Save page as a temporary image file
                with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                    tmp_path = tmp.name
                    pil_img.save(tmp_path, format="PNG")

                try:
                    result = _adaptive_ocr(tmp_path, pil_img=pil_img)

                    if result['text']:
                        all_text_parts.append(result['text'])

                        for block in result['blocks']:
                            block['page'] = page.page_number
                            all_blocks.append(block)
                            total_confidence += block['confidence']
                            total_detections += 1

                except Exception as page_err:
                    print(f"Page {page.page_number} OCR failed: {page_err}")
                finally:
                    try:
                        os.unlink(tmp_path)
                    except OSError:
                        pass

                # Add page separator
                all_text_parts.append(f"\n--- Page {page.page_number} ---\n")

        avg_confidence = total_confidence / total_detections if total_detections else 0.0

        return {
            "raw_text": "\n".join(all_text_parts),
            "blocks": all_blocks,
            "confidence_avg": round(avg_confidence, 3),
            "page_count": page_count,
        }

    except Exception as e:
        return {
            "raw_text": f"Error processing scanned PDF: {str(e)}",
            "blocks": [],
            "confidence_avg": 0.0,
            "page_count": 0,
        }