src/knowledge_parsing/paddle_ocr.py

"""PaddleOCR-VL hosted job API (backend="paddleocr"). Added 2026-09-24.

Kept out of parse.py for the same reason as mistral_ocr.py: the API client stays
separate from caching and normalisation. No GPU — Baidu AI Studio hosts the model.

The flow is the official SDK's (PaddlePaddle/PaddleOCR `api_sdk/go`), read from
source because the job API has no published HTTP reference:

    POST {job_url}              multipart: file, model, optionalPayload   -> data.jobId
    GET  {job_url}/{jobId}      poll until data.state is "done" | "failed"
    GET  data.resultUrl.jsonUrl JSONL, one {"result": {"layoutParsingResults": [...]}}
                                per line, pages in order

Every response is enveloped `{code, msg, data}`; success is `code == 0`.

**Figures are downloaded here, at parse time.** The job returns each figure as a
URL on Baidu object storage, not as bytes. They are fetched immediately and stored
base64 beside the page (`markdown.images_base64`), so the cached response is
self-contained, the normaliser stays offline, and nothing depends on whether those
URLs expire. Page renders (`inputImage`, `outputImages`) are not fetched — nothing
uses them.

PDF and images only: Baidu's API reference lists no other input type.
"""

from __future__ import annotations

import base64
import json
import logging
import mimetypes
import os
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any

import httpx

from .config import PipelineConfig

logger = logging.getLogger("knowledge_parsing.paddle_ocr")

SUPPORTED_SUFFIXES = {".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}

# Explicit rather than server defaults, so a change on Baidu's side cannot quietly
# change what we parse. Chart recognition would turn charts into tables; figures
# are kept as images here, like the other backends.
OPTIONAL_PAYLOAD = {
    "useDocOrientationClassify": False,
    "useDocUnwarping": False,
    "useChartRecognition": False,
}

# The SDK's own polling: 3 s, x1.5, capped at 15 s.
_POLL_FIRST_S = 3.0
_POLL_FACTOR = 1.5
_POLL_MAX_S = 15.0
# A poll that fails on the network is retried this many times in a row before the
# job is given up — one dropped connection should not fail a 60-page document.
_POLL_NETWORK_RETRIES = 3

# Submit retries only for "busy" answers: HTTP 429/503/504 or code 10010 (queue
# full). Spaced out, never an immediate burst at a service that said slow down.
_SUBMIT_RETRY_DELAYS_S = (5.0, 15.0, 45.0)
_QUEUE_FULL_CODE = 10010
_RETRYABLE_STATUS = {429, 503, 504}

_TIMEOUT = httpx.Timeout(120.0, connect=15.0)


class PaddleOCRError(RuntimeError):
    """A short, readable failure — it lands on the document row as-is."""

    def __init__(self, message: str, *, retryable: bool = False) -> None:
        super().__init__(message)
        self.retryable = retryable


def _token_from_settings() -> str | None:
    """The app's settings object, when running inside the service. Never raises."""
    try:
        from src.config.settings import settings
    except Exception:  # noqa: BLE001 - standalone use has no app settings
        return None
    return getattr(settings, "paddleocr_access_token", "") or None


def _token() -> str:
    key = _token_from_settings() or os.environ.get("PADDLEOCR_ACCESS_TOKEN")
    if not key:
        try:
            from dotenv import load_dotenv

            load_dotenv()
            key = os.environ.get("PADDLEOCR_ACCESS_TOKEN")
        except ImportError:
            pass
    if not key:
        raise PaddleOCRError(
            "PADDLEOCR_ACCESS_TOKEN is not set — add it to the app settings / "
            "HuggingFace Space secret, or the repo .env"
        )
    return key


def _data(resp: httpx.Response) -> dict[str, Any]:
    """Unwrap `{code, msg, data}`; raise a short error for anything else."""
    if resp.status_code in (401, 403):
        raise PaddleOCRError(
            f"PaddleOCR rejected the token (HTTP {resp.status_code}) — "
            "check PADDLEOCR_ACCESS_TOKEN"
        )
    try:
        body = resp.json()
    except ValueError:
        raise PaddleOCRError(
            f"PaddleOCR returned HTTP {resp.status_code} with a non-JSON body",
            retryable=resp.status_code in _RETRYABLE_STATUS,
        ) from None
    code = body.get("code") if isinstance(body, dict) else None
    if resp.status_code >= 400 or code not in (0, None):
        raise PaddleOCRError(
            f"PaddleOCR error {code if code not in (None, 0) else resp.status_code}: "
            f"{(body.get('msg') if isinstance(body, dict) else None) or 'no message'}",
            retryable=resp.status_code in _RETRYABLE_STATUS or code == _QUEUE_FULL_CODE,
        )
    data = body.get("data")
    if not isinstance(data, dict):
        raise PaddleOCRError("PaddleOCR response carried no data")
    return data


def _submit(
    client: httpx.Client, source: Path, cfg: PipelineConfig, headers: dict[str, str],
    sleep: Callable[[float], None],
) -> str:
    mime = mimetypes.guess_type(source.name)[0] or "application/octet-stream"
    for attempt in range(len(_SUBMIT_RETRY_DELAYS_S) + 1):
        with source.open("rb") as f:
            resp = client.post(
                cfg.paddle_job_url,
                headers=headers,
                data={"model": cfg.paddle_model, "optionalPayload": json.dumps(OPTIONAL_PAYLOAD)},
                files={"file": (source.name, f, mime)},
            )
        try:
            data = _data(resp)
        except PaddleOCRError as e:
            if e.retryable and attempt < len(_SUBMIT_RETRY_DELAYS_S):
                delay = _SUBMIT_RETRY_DELAYS_S[attempt]
                logger.warning("paddleocr submit busy (%s); retrying in %.0fs", e, delay)
                sleep(delay)
                continue
            raise
        job_id = data.get("jobId")
        if not job_id:
            raise PaddleOCRError("PaddleOCR accepted the file but returned no jobId")
        return str(job_id)
    raise PaddleOCRError("PaddleOCR stayed busy; the file was not submitted")  # unreachable


def _wait_for_result(
    client: httpx.Client, job_id: str, cfg: PipelineConfig, headers: dict[str, str],
    sleep: Callable[[float], None], clock: Callable[[], float],
) -> str:
    """Poll until the job is done; return its `jsonUrl`."""
    deadline = clock() + cfg.paddle_poll_timeout_s
    interval = _POLL_FIRST_S
    network_failures = 0
    state = "submitted"
    while True:
        try:
            data = _data(client.get(f"{cfg.paddle_job_url}/{job_id}", headers=headers))
            network_failures = 0
        except httpx.TransportError as e:
            network_failures += 1
            if network_failures > _POLL_NETWORK_RETRIES:
                raise PaddleOCRError(f"PaddleOCR job {job_id}: status unreachable ({e!r})") from e
            data = {}
        state = data.get("state") or state
        if state == "done":
            url = (data.get("resultUrl") or {}).get("jsonUrl")
            if not url:
                raise PaddleOCRError(f"PaddleOCR job {job_id} finished without a result URL")
            return str(url)
        if state == "failed":
            raise PaddleOCRError(
                f"PaddleOCR job {job_id} failed: {data.get('errorMsg') or 'no reason given'}"
            )
        if clock() >= deadline:
            raise PaddleOCRError(
                f"PaddleOCR job {job_id} still '{state}' after {cfg.paddle_poll_timeout_s:.0f}s"
            )
        progress = data.get("extractProgress") or {}
        logger.info(
            "paddleocr job %s %s (%s/%s pages)", job_id, state,
            progress.get("extractedPages", "?"), progress.get("totalPages", "?"),
        )
        sleep(interval)
        interval = min(interval * _POLL_FACTOR, _POLL_MAX_S)


def _result_pages(jsonl: str) -> list[dict[str, Any]]:
    """Flatten the JSONL into one list of page results, in order."""
    pages: list[dict[str, Any]] = []
    for line in jsonl.splitlines():
        if not line.strip():
            continue
        result = json.loads(line).get("result") or {}
        pages.extend(result.get("layoutParsingResults") or [])
    return pages


def _embed_figures(client: httpx.Client, pages: list[dict[str, Any]]) -> None:
    """Fetch each page's figure URLs into `markdown.images_base64` (in place).

    Never-throw per figure: one failed download costs that figure its bytes (the
    normaliser still references it), never the document. The token is NOT sent —
    these are pre-signed storage URLs on another host.
    """
    for page in pages:
        markdown = page.get("markdown") or {}
        stored: dict[str, str] = {}
        for path, ref in (markdown.get("images") or {}).items():
            if not isinstance(ref, str) or not ref:
                continue
            if not ref.startswith(("http://", "https://")):
                stored[path] = ref           # already base64 (the sync API's shape)
                continue
            try:
                r = client.get(ref)
                r.raise_for_status()
                stored[path] = base64.b64encode(r.content).decode("ascii")
            except Exception as e:  # noqa: BLE001 - one figure is not worth a document
                logger.warning("paddleocr figure download failed for %s: %r", path, e)
        markdown["images_base64"] = stored
        page["markdown"] = markdown


def run_ocr(
    source: Path,
    cfg: PipelineConfig,
    *,
    client: httpx.Client | None = None,
    sleep: Callable[[float], None] = time.sleep,
    clock: Callable[[], float] = time.monotonic,
) -> dict[str, Any]:
    """Parse one PDF/image with PaddleOCR-VL. Returns `{model, job_id, pages}`.

    `client`, `sleep` and `clock` are injectable so the whole flow is testable
    without the network.
    """
    if source.suffix.lower() not in SUPPORTED_SUFFIXES:
        raise PaddleOCRError(
            f"PaddleOCR takes PDF or images only, not {source.suffix or 'this file type'}"
        )
    headers = {"Authorization": f"Bearer {_token()}"}
    owns_client = client is None
    http = client or httpx.Client(timeout=_TIMEOUT, follow_redirects=True)
    try:
        job_id = _submit(http, source, cfg, headers, sleep)
        json_url = _wait_for_result(http, job_id, cfg, headers, sleep, clock)
        resp = http.get(json_url)
        resp.raise_for_status()
        pages = _result_pages(resp.text)
        if not pages:
            raise PaddleOCRError(f"PaddleOCR job {job_id} returned no pages")
        _embed_figures(http, pages)
        return {"model": cfg.paddle_model, "job_id": job_id, "pages": pages}
    finally:
        if owns_client:
            http.close()