src/services/ocr_engine.py

"""
ocr_engine.py
─────────────
PaddleOCR inference service with:
  - Angle classifier enabled (handles rotated/upside-down text lines)
  - Confidence-aware result passing (raw boxes + scores to the formatter)
  - Tesseract fallback for low-confidence pages
"""

import os
import asyncio
from typing import Dict, List, Optional, Tuple
import numpy as np

# Disable GPU and CPU MKL-DNN optimizations to prevent container segmentation faults
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_use_xdnn"] = "0"
os.environ["OMP_NUM_THREADS"] = "1"

from paddleocr import PaddleOCR
from src.config import settings

# Threshold: if average page confidence falls below this, run Tesseract fallback
TESS_FALLBACK_THRESHOLD = 0.65


class OCREngineService:
    def __init__(self):
        # In-memory cache of instantiated PaddleOCR engines keyed by language
        self._engines: Dict[str, PaddleOCR] = {}
        # Mutex lock to serialize CPU inference and avoid resource thrashing
        self._lock = asyncio.Lock()

    def get_engine(self, lang: str) -> PaddleOCR:
        """Retrieves or initialises a PaddleOCR engine for the given language."""
        lang = lang.lower().strip()

        if lang not in settings.SUPPORTED_LANGUAGES:
            lang = settings.DEFAULT_LANGUAGE

        if lang not in self._engines:
            self._engines[lang] = PaddleOCR(
                use_angle_cls=True,          # Detects 0°/180° text orientation (fixes upside-down lines)
                lang=lang,
                show_log=False,
                enable_mkldnn=False,
                det_limit_side_len=3000,     # Preserve high-res details for tiny text (default 960)
                det_db_box_thresh=0.4,       # Lower threshold — detect faint/thin boxes (was 0.5)
                det_db_score_mode="slow",    # Accurate boundary evaluation (default 'fast')
                rec_db_score_mode="slow",
                use_space_char=True,         # Recognise spaces between words
                cls_thresh=0.9,              # High bar for angle flip — only flip when very confident
            )
        return self._engines[lang]

    async def _detect_language(self, img_np: np.ndarray, loop) -> str:
        """
        Uses Tesseract OSD to detect the writing system/script, and maps it to a supported language.
        """
        try:
            import pytesseract
            import cv2
            import re

            # Tesseract works best on grayscale
            gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)

            osd_data = await loop.run_in_executor(
                None,
                lambda: pytesseract.image_to_osd(gray, config="--psm 0")
            )

            # Parse Script Name (e.g. 'Script: Latin', 'Script: Arabic', etc.)
            match = re.search(r"Script:\s+(\w+)", osd_data)
            if match:
                script = match.group(1).lower().strip()

                # Map script to PaddleOCR language code
                script_map = {
                    "latin": "en",       # Latin alphabet covers English, Spanish, Portuguese, French, German, Italian
                    "devanagari": "hi",  # Hindi
                    "arabic": "ar",      # Arabic
                    "han": "ch",         # Chinese (Simplified)
                    "cyrillic": "ru",    # Russian
                    "telugu": "te",      # Telugu
                    "tamil": "ta",       # Tamil
                    "japanese": "japan", # Japanese
                    "korean": "korean"   # Korean
                }

                detected = script_map.get(script, "en")
                print(f"OSD detected script '{script}', mapped to language code '{detected}'")
                return detected

        except Exception as e:
            print(f"OSD language detection failed: {e}. Falling back to default 'en'.")

        return "en"

    async def run_ocr_raw(self, img_np: np.ndarray, lang: str):
        """
        Runs PaddleOCR and returns the raw result list (bounding boxes + text + confidence).
        The caller (routes.py) passes this to the markdown formatter.
        """
        async with self._lock:
            loop = asyncio.get_running_loop()

            # Resolve auto language detection if requested
            if lang.lower().strip() in ("auto", "detect"):
                resolved_lang = await self._detect_language(img_np, loop)
            else:
                resolved_lang = lang

            ocr  = self.get_engine(resolved_lang)

            raw = await loop.run_in_executor(
                None,
                lambda: ocr.ocr(img_np, cls=True)
            )

        # raw is [[result_page1], ...] — we only handle single-page images
        page = raw[0] if raw else []

        # Optionally run Tesseract fallback if average confidence is low
        if _should_run_tesseract(page):
            tess_lines = await _run_tesseract_fallback(img_np, loop)
            page = _merge_tesseract_lines(page, tess_lines)

        return page


# ─────────────────────────────────────────────
# Tesseract fallback helpers
# ─────────────────────────────────────────────

def _should_run_tesseract(page_result) -> bool:
    """
    Returns True if the average PaddleOCR confidence score for the page
    is below TESS_FALLBACK_THRESHOLD, indicating the page may have zones
    of text that PaddleOCR struggled with.
    """
    if not page_result:
        return True  # Empty result — definitely try fallback

    scores = [item[1][1] for item in page_result if item and item[1]]
    if not scores:
        return True

    avg_conf = sum(scores) / len(scores)
    return avg_conf < TESS_FALLBACK_THRESHOLD


async def _run_tesseract_fallback(img_np: np.ndarray, loop) -> List[str]:
    """
    Runs pytesseract in sparse-text mode as a fallback for difficult pages.
    psm 11 = Sparse text. Find as much text as possible in no particular order.
    oem 1  = LSTM neural net engine only.
    """
    try:
        import pytesseract
        import cv2

        # Tesseract works best on grayscale
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)

        config = "--oem 1 --psm 11"

        text = await loop.run_in_executor(
            None,
            lambda: pytesseract.image_to_string(gray, config=config, lang="eng")
        )

        # Split into lines, strip whitespace, drop empty
        lines = [l.strip() for l in text.splitlines() if l.strip()]
        return lines

    except ImportError:
        # pytesseract not installed — silently skip
        return []
    except Exception as e:
        print(f"Tesseract fallback failed: {e}")
        return []


def _merge_tesseract_lines(paddle_page, tess_lines: List[str]):
    """
    Merge Tesseract-recovered lines into the PaddleOCR result.

    Strategy: collect all text already captured by PaddleOCR, then
    append any Tesseract lines that are not already present as a
    'synthetic' low-priority item (no bounding box, confidence 0.60).
    This ensures the formatter sees them without duplicating content.
    """
    if not tess_lines:
        return paddle_page

    existing_texts = set()
    if paddle_page:
        for item in paddle_page:
            if item and item[1]:
                existing_texts.add(item[1][0].strip().lower())

    merged = list(paddle_page) if paddle_page else []

    for line in tess_lines:
        normalised = line.strip().lower()
        if normalised and normalised not in existing_texts:
            # Synthetic entry: no real bounding box — place at the bottom of the page
            # The formatter handles this gracefully (falls back to plain text for lines without boxes)
            merged.append([
                [[0, 9999], [100, 9999], [100, 10000], [0, 10000]],
                (line, 0.60)
            ])
            existing_texts.add(normalised)

    return merged


# Global singleton
ocr_service = OCREngineService()