ocr/paddle_engine.py

"""
InvoiceForge AI — ocr/paddle_engine.py
PaddleOCR 3.x + PPStructureV3 engine wrapper.

Provides:
  - PaddleEngine: singleton OCR engine with lazy initialisation
  - run_ocr(): returns standardised token list
  - run_structure(): PPStructureV3 table + layout analysis
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any

import cv2
import numpy as np

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# DATA CLASSES
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class OCRToken:
    """Single recognised text token from any OCR engine."""
    text: str
    confidence: float
    bbox: list  # [[x1,y1],[x2,y1],[x2,y2],[x1,y2]]
    x: float = 0.0
    y: float = 0.0
    engine: str = "paddleocr"

    def __post_init__(self) -> None:
        if self.bbox:
            xs = [p[0] for p in self.bbox]
            ys = [p[1] for p in self.bbox]
            self.x = float(np.mean(xs))
            self.y = float(np.mean(ys))


@dataclass
class TableCell:
    """Single cell from PPStructureV3 table extraction."""
    row: int
    col: int
    row_span: int
    col_span: int
    text: str
    bbox: list


@dataclass
class TableResult:
    """Full table extracted by PPStructureV3."""
    cells: list[TableCell] = field(default_factory=list)
    html: str = ""
    confidence: float = 0.0


# ─────────────────────────────────────────────────────────────────────────────
# PADDLE ENGINE SINGLETON
# ─────────────────────────────────────────────────────────────────────────────

class PaddleEngine:
    """
    Thread-safe singleton wrapper around PaddleOCR and PPStructure.

    Initialization is deferred until first use to avoid penalising startup.
    Usage:
        engine = PaddleEngine.instance()
        tokens = engine.run_ocr(img_array)
    """

    _ocr_instance: Any = None
    _struct_instance: Any = None

    @classmethod
    def instance(cls) -> "PaddleEngine":
        """Return the singleton PaddleEngine."""
        obj = cls.__new__(cls)
        return obj

    def _get_ocr(self) -> Any:
        """Lazily initialise PaddleOCR."""
        if PaddleEngine._ocr_instance is None:
            logger.info("Initialising PaddleOCR engine …")
            try:
                from paddleocr import PaddleOCR  # type: ignore
                PaddleEngine._ocr_instance = PaddleOCR(
                    use_angle_cls=True,
                    lang="en",
                    show_log=False,
                    use_gpu=False,
                    # v3 models — highest accuracy for English invoices
                    det_model_dir=None,
                    rec_model_dir=None,
                    cls_model_dir=None,
                )
                logger.info("PaddleOCR initialised successfully.")
            except Exception as exc:
                logger.error("PaddleOCR init failed: %s", exc)
                raise
        return PaddleEngine._ocr_instance

    def _get_structure(self) -> Any:
        """Lazily initialise PPStructureV3."""
        if PaddleEngine._struct_instance is None:
            logger.info("Initialising PPStructure engine …")
            try:
                from paddleocr import PPStructure  # type: ignore
                PaddleEngine._struct_instance = PPStructure(
                    table=True,
                    ocr=True,
                    show_log=False,
                    use_gpu=False,
                )
                logger.info("PPStructure initialised successfully.")
            except Exception as exc:
                logger.error("PPStructure init failed: %s", exc)
                raise
        return PaddleEngine._struct_instance

    # ── OCR ───────────────────────────────────────────────────────────────────

    def run_ocr(self, img: np.ndarray) -> list[OCRToken]:
        """
        Run PaddleOCR on a preprocessed image.

        Args:
            img: Grayscale or BGR numpy array.

        Returns:
            List of OCRToken objects sorted by (y, x).
        """
        ocr = self._get_ocr()
        # PaddleOCR accepts BGR or gray; ensure uint8
        if img.dtype != np.uint8:
            img = img.astype(np.uint8)

        result = ocr.ocr(img, cls=True)
        tokens: list[OCRToken] = []

        if not result or not result[0]:
            return tokens

        for line in result[0]:
            if line is None:
                continue
            bbox, (text, conf) = line[0], line[1]
            if conf < 0.25:  # discard near-noise detections
                continue
            tokens.append(
                OCRToken(
                    text=text.strip(),
                    confidence=float(conf),
                    bbox=bbox,
                    engine="paddleocr",
                )
            )

        tokens.sort(key=lambda t: (t.y, t.x))
        logger.debug("PaddleOCR: %d tokens extracted.", len(tokens))
        return tokens

    # ── STRUCTURE / TABLE ─────────────────────────────────────────────────────

    def run_structure(self, img_bgr: np.ndarray) -> list[dict]:
        """
        Run PPStructureV3 layout + table analysis.

        Args:
            img_bgr: BGR numpy array (colour input required by PPStructure).

        Returns:
            List of region dicts from PPStructure, each containing:
              type, bbox, res (OCR or table HTML), and score.
        """
        if len(img_bgr.shape) == 2:
            img_bgr = cv2.cvtColor(img_bgr, cv2.COLOR_GRAY2BGR)

        struct = self._get_structure()
        try:
            result = struct(img_bgr)
            logger.debug("PPStructure: %d regions detected.", len(result))
            return result or []
        except Exception as exc:
            logger.warning("PPStructure failed: %s", exc)
            return []

    def extract_tables(self, img_bgr: np.ndarray) -> list[TableResult]:
        """
        Extract all tables from the image using PPStructureV3.
        Parses the HTML output into structured TableCell objects.

        Returns:
            List of TableResult objects (one per detected table).
        """
        regions = self.run_structure(img_bgr)
        tables: list[TableResult] = []

        for region in regions:
            if region.get("type", "").lower() != "table":
                continue
            res = region.get("res", {})
            html = res.get("html", "") if isinstance(res, dict) else ""
            score = float(region.get("score", 0.0))
            cells = _parse_table_html(html)
            tables.append(TableResult(cells=cells, html=html, confidence=score))

        return tables


# ─────────────────────────────────────────────────────────────────────────────
# TABLE HTML PARSER
# ─────────────────────────────────────────────────────────────────────────────

def _parse_table_html(html: str) -> list[TableCell]:
    """
    Parse HTML table string from PPStructureV3 into TableCell objects.
    Handles rowspan and colspan attributes.
    """
    if not html:
        return []

    try:
        import re

        cells: list[TableCell] = []
        row_idx = 0

        # Find all table rows
        tr_pattern = re.compile(r"<tr[^>]*>(.*?)</tr>", re.DOTALL | re.IGNORECASE)
        td_pattern = re.compile(
            r"<td([^>]*)>(.*?)</td>", re.DOTALL | re.IGNORECASE
        )
        attr_pattern = re.compile(r'(\w+)=["\'](\d+)["\']')

        for tr_match in tr_pattern.finditer(html):
            row_html = tr_match.group(1)
            col_idx = 0
            for td_match in td_pattern.finditer(row_html):
                attrs_str = td_match.group(1)
                content = re.sub(r"<[^>]+>", "", td_match.group(2)).strip()

                attrs: dict[str, int] = {}
                for attr_m in attr_pattern.finditer(attrs_str):
                    attrs[attr_m.group(1).lower()] = int(attr_m.group(2))

                row_span = attrs.get("rowspan", 1)
                col_span = attrs.get("colspan", 1)

                cells.append(
                    TableCell(
                        row=row_idx,
                        col=col_idx,
                        row_span=row_span,
                        col_span=col_span,
                        text=content,
                        bbox=[],
                    )
                )
                col_idx += col_span
            row_idx += 1

        return cells

    except Exception as exc:
        logger.warning("Table HTML parse failed: %s", exc)
        return []


# ─────────────────────────────────────────────────────────────────────────────
# MODULE-LEVEL SINGLETON
# ─────────────────────────────────────────────────────────────────────────────

_engine: PaddleEngine | None = None


def get_paddle_engine() -> PaddleEngine:
    """Return the module-level PaddleEngine singleton."""
    global _engine
    if _engine is None:
        _engine = PaddleEngine.instance()
    return _engine