ragflow/deepdoc/parser/opendataloader_parser.py

from __future__ import annotations

import html
import json
import logging
import os
import re
from dataclasses import dataclass
from enum import Enum
from io import BytesIO
from os import PathLike
from pathlib import Path
from typing import Any, Callable, Iterable, Optional

import pdfplumber
import requests
from PIL import Image

from common.constants import MAXIMUM_PAGE_NUMBER

try:
    from deepdoc.parser.pdf_parser import RAGFlowPdfParser
except Exception:

    class RAGFlowPdfParser:
        pass


from deepdoc.parser.utils import extract_pdf_outlines


class OpenDataLoaderContentType(str, Enum):
    IMAGE = "image"
    TABLE = "table"
    TEXT = "text"
    EQUATION = "equation"


@dataclass
class _BBox:
    page_no: int
    x0: float
    y0: float
    x1: float
    y1: float


_TEXT_TYPES = {"heading", "title", "paragraph", "text", "list", "list_item", "caption"}
_TABLE_TYPES = {"table"}
_IMAGE_TYPES = {"image", "picture", "figure"}
_FORMULA_TYPES = {"formula", "equation"}


def _as_float(v) -> Optional[float]:
    try:
        return float(v)
    except Exception:
        return None


def _bbox_from_element(el: dict) -> Optional[_BBox]:
    bb = el.get("bounding box") or el.get("bounding_box") or el.get("bbox")
    pn = el.get("page number")
    if pn is None:
        pn = el.get("page_number")
    if pn is None:
        pn = el.get("page")
    if bb is None or pn is None:
        return None
    if not isinstance(bb, (list, tuple)) or len(bb) < 4:
        return None
    coords = [_as_float(x) for x in bb[:4]]
    if any(c is None for c in coords):
        return None
    try:
        page_no = int(pn)
    except Exception:
        return None
    # OpenDataLoader emits [left, bottom, right, top] in PDF points.
    left, bottom, right, top = coords
    x0, x1 = min(left, right), max(left, right)
    y0, y1 = min(bottom, top), max(bottom, top)
    return _BBox(page_no=page_no, x0=x0, y0=y0, x1=x1, y1=y1)


def _iter_elements(node: Any) -> Iterable[dict]:
    if isinstance(node, dict):
        if "type" in node and ("content" in node or "text" in node or "cells" in node or "table_cells" in node):
            yield node
        for v in node.values():
            yield from _iter_elements(v)
    elif isinstance(node, list):
        for item in node:
            yield from _iter_elements(item)


def _element_text(el: dict) -> str:
    content = el.get("content")
    if isinstance(content, str):
        return content
    text = el.get("text")
    if isinstance(text, str):
        return text
    # tables may expose cells; join row-wise if needed
    cells = el.get("cells")
    if isinstance(cells, list):
        rows: dict[int, list[str]] = {}
        for c in cells:
            if not isinstance(c, dict):
                continue
            row = c.get("row") or c.get("row_index") or 0
            rows.setdefault(int(row), []).append(str(c.get("content") or c.get("text") or ""))
        return "\n".join(" | ".join(v) for _, v in sorted(rows.items()))
    return ""


def _element_html(el: dict) -> str:
    for key in ("html", "html_content"):
        v = el.get(key)
        if isinstance(v, str) and v.strip():
            return v
    # Docling tables: rebuild HTML from table_cells
    cells = el.get("table_cells")
    if isinstance(cells, list) and cells:
        num_rows = el.get("num_rows") or 0
        num_cols = el.get("num_cols") or 0
        if num_rows > 0 and num_cols > 0:
            grid: list[list[dict]] = [[{} for _ in range(num_cols)] for _ in range(num_rows)]
            for c in cells:
                if not isinstance(c, dict):
                    continue
                r = c.get("start_row_offset_idx", 0)
                col = c.get("start_col_offset_idx", 0)
                row_span = max(1, c.get("row_span", 1) or 1)
                col_span = max(1, c.get("col_span", 1) or 1)
                text = html.escape(c.get("text", ""))
                if 0 <= r < num_rows and 0 <= col < num_cols:
                    grid[r][col] = {"text": text, "row_span": row_span, "col_span": col_span}
                    # Mark spanned positions as occupied so they are skipped
                    for dr in range(row_span):
                        for dc in range(col_span):
                            if dr == 0 and dc == 0:
                                continue
                            rr, cc = r + dr, col + dc
                            if 0 <= rr < num_rows and 0 <= cc < num_cols:
                                grid[rr][cc] = {"_skip": True}
            rows_html = []
            for row in grid:
                cols = []
                for cell_data in row:
                    if cell_data.get("_skip"):
                        continue
                    if not cell_data:
                        continue
                    attrs = ""
                    if cell_data.get("row_span", 1) > 1:
                        attrs += f' rowspan="{cell_data["row_span"]}"'
                    if cell_data.get("col_span", 1) > 1:
                        attrs += f' colspan="{cell_data["col_span"]}"'
                    cols.append(f"<td{attrs}>{cell_data['text']}</td>")
                rows_html.append("<tr>" + "".join(cols) + "</tr>")
            return "<table>" + "".join(rows_html) + "</table>"
    return ""


# ---------------------------------------------------------------------------
# DoclingDocument → intermediate format adapter
# ---------------------------------------------------------------------------
# The opendataloader-pdf-hybrid service (/v1/convert/file) returns a
# DoclingDocument JSON whose structure differs from the original OpenDataLoader
# format.  This adapter normalises it so _transfer_from_json can consume it.
# ---------------------------------------------------------------------------

_DOCLING_TEXT_LABELS = {
    "paragraph",
    "text",
    "section_header",
    "title",
    "heading",
    "list_item",
    "code",
    "caption",
    "footnote",
    "page_header",
    "page_footer",
}
_DOCLING_TABLE_LABELS = {"table"}
_DOCLING_PICTURE_LABELS = {"picture", "figure"}
_DOCLING_FORMULA_LABELS = {"formula"}


def _normalize_docling_response(doc: dict, page_heights: dict[int, float] | None = None) -> list[dict]:
    """Convert a DoclingDocument JSON into the intermediate element format
    that _iter_elements / _transfer_from_json / _bbox_from_element expect.

    Each returned element dict has:
      - "type": element type string (table, picture, paragraph, section_header, …)
      - "text": or "content" for text content
      - "page_number": page number
      - "bounding_box": [left, bottom, right, top] flat array in PDF points
      - "html": rebuilt HTML for tables (optional)
      - "cells": / "table_cells": for tables (optional)

    Parameters:
      doc: the DoclingDocument JSON dict.
      page_heights: optional mapping from page_no (1-based) to page height
                    in PDF points.  When provided, Y coordinates from
                    Docling's TOPLEFT origin are flipped to bottom-left
                    origin.  When absent the raw values are used (may
                    produce incorrect positions).
    """
    if page_heights is None:
        page_heights = {}
    elements: list[dict] = []

    # --- texts ---
    for item in doc.get("texts") or []:
        if not isinstance(item, dict):
            continue
        lbl = (item.get("label") or "").lower()
        if lbl in {"page_header", "page_footer"}:
            continue
        el_type = lbl  # preserve original label (section_header, paragraph, etc.)
        text = item.get("text", "")
        page_no, bbox_arr = _extract_docling_prov(item, page_heights)
        el: dict[str, Any] = {"type": el_type, "text": text}
        if page_no is not None:
            el["page_number"] = page_no
        if bbox_arr is not None:
            el["bounding_box"] = bbox_arr
        elements.append(el)

    # --- tables ---
    for item in doc.get("tables") or []:
        if not isinstance(item, dict):
            continue
        el_type = "table"
        page_no, bbox_arr = _extract_docling_prov(item, page_heights)
        el: dict[str, Any] = {"type": el_type, "table_cells": item.get("table_cells", []), "num_rows": item.get("num_rows", 0), "num_cols": item.get("num_cols", 0)}
        if page_no is not None:
            el["page_number"] = page_no
        if bbox_arr is not None:
            el["bounding_box"] = bbox_arr
        elements.append(el)

    # --- pictures ---
    for item in doc.get("pictures") or []:
        if not isinstance(item, dict):
            continue
        el_type = "picture"
        # Try to get caption from embedded data or referenced captions
        caption = ""
        # Some picture items may have a "data" field with description
        data = item.get("data")
        if isinstance(data, dict):
            caption = data.get("description", "")
        page_no, bbox_arr = _extract_docling_prov(item, page_heights)
        el: dict[str, Any] = {"type": el_type, "content": caption}
        if page_no is not None:
            el["page_number"] = page_no
        if bbox_arr is not None:
            el["bounding_box"] = bbox_arr
        elements.append(el)

    return elements


def _extract_docling_prov(item: dict, page_heights: dict[int, float] | None = None) -> tuple[Optional[int], Optional[list[float]]]:
    """Extract page number and bounding box from a Docling item's prov list.

    Docling prov format: [{"page_no": int, "bbox": {"l": float, "t": float,
    "r": float, "b": float, "coord_origin": "TOPLEFT"}}]

    Returns (page_no, [left, bottom, right, top]) in PDF points with
    bottom-left origin (matching OpenDataLoader convention).

    When a page_heights mapping is provided, Docling's TOPLEFT Y coordinates
    are flipped using the page height.  Without it the raw values are
    returned as-is (which places the bounding box at the wrong vertical
    position).
    """
    if page_heights is None:
        page_heights = {}
    prov_list = item.get("prov")
    if not isinstance(prov_list, list) or not prov_list:
        return None, None
    # Use the first provenance entry
    prov = prov_list[0]
    page_no = prov.get("page_no")
    if page_no is None:
        return None, None
    bbox = prov.get("bbox")
    if not isinstance(bbox, dict):
        return None, None
    left = _as_float(bbox.get("l"))
    top_from_top = _as_float(bbox.get("t"))
    right = _as_float(bbox.get("r"))
    bottom_from_top = _as_float(bbox.get("b"))
    if any(v is None for v in (left, top_from_top, right, bottom_from_top)):
        return None, None
    # Docling uses top-left origin (t = distance from top of page).
    # Convert to bottom-left origin using page height when available.
    page_height = page_heights.get(int(page_no))
    if page_height is not None:
        bottom_to_bottom = page_height - bottom_from_top
        top_to_bottom = page_height - top_from_top
    else:
        # Fallback: raw values (incorrect positions but better than nothing)
        bottom_to_bottom = bottom_from_top
        top_to_bottom = top_from_top
    return int(page_no), [left, bottom_to_bottom, right, top_to_bottom]


class OpenDataLoaderParser(RAGFlowPdfParser):
    def __init__(self):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.page_images: list[Image.Image] = []
        self.page_from = 0
        self.page_to = 10_000
        self.outlines = []
        self.api_url = os.environ.get("OPENDATALOADER_APISERVER", "").rstrip("/")
        self.api_key = os.environ.get("OPENDATALOADER_API_KEY", "").strip()
        try:
            self.timeout = int(os.environ.get("OPENDATALOADER_TIMEOUT", "600") or "600")
        except ValueError:
            self.logger.warning("[OpenDataLoader] Invalid OPENDATALOADER_TIMEOUT, falling back to 600s")
            self.timeout = 600

    def check_installation(self) -> bool:
        """Return True when the OpenDataLoader service is reachable."""
        if not self.api_url:
            self.logger.warning("[OpenDataLoader] OPENDATALOADER_APISERVER is not set. Start the opendataloader service and set the env var.")
            return False
        try:
            headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
            resp = requests.get(f"{self.api_url}/health", timeout=5, headers=headers)
            if resp.status_code == 200:
                return True
            self.logger.warning(f"[OpenDataLoader] Health check returned {resp.status_code}: {resp.text[:200]}")
            return False
        except Exception as exc:
            self.logger.warning(f"[OpenDataLoader] Health check failed: {exc}")
            return False

    def __images__(self, fnm, zoomin: int = 1, page_from=0, page_to=MAXIMUM_PAGE_NUMBER, callback=None):
        self.page_from = page_from
        self.page_to = page_to
        bytes_io = None
        try:
            if not isinstance(fnm, (str, PathLike)):
                bytes_io = fnm if isinstance(fnm, BytesIO) else BytesIO(fnm)
            opener = pdfplumber.open(fnm) if isinstance(fnm, (str, PathLike)) else pdfplumber.open(bytes_io)
            with opener as pdf:
                pages = pdf.pages[page_from:page_to]
                self.page_images = [p.to_image(resolution=72 * zoomin, antialias=True).original for p in pages]
        except Exception as e:
            self.page_images = []
            self.logger.exception(e)
        finally:
            if bytes_io:
                bytes_io.close()

    def _make_line_tag(self, bbox: _BBox) -> str:
        if bbox is None:
            return ""
        # Guard: only emit a crop tag when the page was actually rendered.
        if not self.page_images or bbox.page_no <= 0 or len(self.page_images) < bbox.page_no:
            return ""
        x0, x1 = bbox.x0, bbox.x1
        # OpenDataLoader bbox uses PDF coordinate space (origin bottom-left).
        # Convert to image-space (origin top-left) by subtracting from page height.
        _, page_height = self.page_images[bbox.page_no - 1].size
        top = page_height - bbox.y1
        bott = page_height - bbox.y0
        return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##".format(bbox.page_no, x0, x1, top, bott)

    @staticmethod
    def extract_positions(txt: str) -> list[tuple[list[int], float, float, float, float]]:
        poss = []
        for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", txt):
            pn, left, right, top, bottom = tag.strip("#").strip("@").split("\t")
            left, right, top, bottom = float(left), float(right), float(top), float(bottom)
            poss.append(([int(p) - 1 for p in pn.split("-")], left, right, top, bottom))
        return poss

    def crop(self, text: str, ZM: int = 1, need_position: bool = False):
        if not self.page_images:
            return (None, None) if need_position else None
        imgs = []
        poss = self.extract_positions(text)
        if not poss:
            return (None, None) if need_position else None
        # Drop positions whose page indices fall outside the rendered range.
        max_page = len(self.page_images) - 1
        poss = [p for p in poss if all(0 <= pn <= max_page for pn in p[0])]
        if not poss:
            return (None, None) if need_position else None
        GAP = 6
        pos = poss[0]
        poss.insert(0, ([pos[0][0]], pos[1], pos[2], max(0, pos[3] - 120), max(pos[3] - GAP, 0)))
        pos = poss[-1]
        poss.append(([pos[0][-1]], pos[1], pos[2], min(self.page_images[pos[0][-1]].size[1], pos[4] + GAP), min(self.page_images[pos[0][-1]].size[1], pos[4] + 120)))
        positions = []
        for ii, (pns, left, right, top, bottom) in enumerate(poss):
            if bottom <= top:
                bottom = top + 4
            img0 = self.page_images[pns[0]]
            x0, y0, x1, y1 = int(left), int(top), int(right), int(min(bottom, img0.size[1]))
            crop0 = img0.crop((x0, y0, x1, y1))
            imgs.append(crop0)
            if 0 < ii < len(poss) - 1:
                positions.append((pns[0] + self.page_from, x0, x1, y0, y1))
            remain_bottom = bottom - img0.size[1]
            for pn in pns[1:]:
                if remain_bottom <= 0:
                    break
                page = self.page_images[pn]
                x0, y0, x1, y1 = int(left), 0, int(right), int(min(remain_bottom, page.size[1]))
                cimgp = page.crop((x0, y0, x1, y1))
                imgs.append(cimgp)
                if 0 < ii < len(poss) - 1:
                    positions.append((pn + self.page_from, x0, x1, y0, y1))
                remain_bottom -= page.size[1]
        if not imgs:
            return (None, None) if need_position else None
        height = sum(i.size[1] + GAP for i in imgs)
        width = max(i.size[0] for i in imgs)
        pic = Image.new("RGB", (width, int(height)), (245, 245, 245))
        h = 0
        for ii, img in enumerate(imgs):
            if ii == 0 or ii + 1 == len(imgs):
                img = img.convert("RGBA")
                overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
                overlay.putalpha(128)
                img = Image.alpha_composite(img, overlay).convert("RGB")
            pic.paste(img, (0, int(h)))
            h += img.size[1] + GAP
        return (pic, positions) if need_position else pic

    def _cropout_region(self, bbox: _BBox, zoomin: int = 1):
        if not self.page_images:
            return None, ""
        idx = (bbox.page_no - 1) - self.page_from
        if idx < 0 or idx >= len(self.page_images):
            return None, ""
        page_img = self.page_images[idx]
        W, H = page_img.size
        x0 = max(0.0, min(float(bbox.x0), W - 1))
        y0 = max(0.0, min(float(H - bbox.y1), H - 1))
        x1 = max(x0 + 1.0, min(float(bbox.x1), W))
        y1 = max(y0 + 1.0, min(float(H - bbox.y0), H))
        try:
            crop = page_img.crop((int(x0), int(y0), int(x1), int(y1))).convert("RGB")
        except Exception:
            return None, ""
        pos = (bbox.page_no - 1 if bbox.page_no > 0 else 0, x0, x1, y0, y1)
        return crop, [pos]

    def _classify(self, el_type: str) -> str:
        t = (el_type or "").lower()
        if t in _TABLE_TYPES:
            return OpenDataLoaderContentType.TABLE.value
        if t in _IMAGE_TYPES:
            return OpenDataLoaderContentType.IMAGE.value
        if t in _FORMULA_TYPES:
            return OpenDataLoaderContentType.EQUATION.value
        # Preserve the original structural type (heading, title, paragraph,
        # list, caption, …) so downstream parsers can apply heading/title heuristics.
        return t if t else OpenDataLoaderContentType.TEXT.value

    def _transfer_from_json(self, root: Any, parse_method: str):
        sections: list[tuple[str, ...]] = []
        tables: list = []
        for el in _iter_elements(root):
            el_type = self._classify(el.get("type", ""))
            bbox = _bbox_from_element(el)
            tag = self._make_line_tag(bbox) if bbox else ""

            if el_type == OpenDataLoaderContentType.TABLE.value:
                html = _element_html(el) or _element_text(el)
                img = None
                positions = ""
                if bbox:
                    img, positions = self._cropout_region(bbox)
                tables.append(((img, html), positions if positions else ""))
                continue

            if el_type == OpenDataLoaderContentType.IMAGE.value:
                img = None
                positions = ""
                if bbox:
                    img, positions = self._cropout_region(bbox)
                caption = _element_text(el)
                tables.append(((img, [caption] if caption else [""]), positions if positions else ""))
                continue

            text = _element_text(el).strip()
            if not text:
                continue
            if parse_method in {"manual", "pipeline"}:
                sections.append((text, el_type, tag))
            elif parse_method == "paper":
                sections.append((text + tag, el_type))
            else:
                sections.append((text, tag))
        return sections, tables

    @staticmethod
    def _sections_from_markdown(md: str, parse_method: str) -> list[tuple[str, ...]]:
        txt = (md or "").strip()
        if not txt:
            return []
        if parse_method in {"manual", "pipeline"}:
            return [(txt, OpenDataLoaderContentType.TEXT.value, "")]
        if parse_method == "paper":
            return [(txt, OpenDataLoaderContentType.TEXT.value)]
        return [(txt, "")]

    def parse_pdf(
        self,
        filepath: str | PathLike[str],
        binary: BytesIO | bytes | None = None,
        callback: Optional[Callable] = None,
        *,
        parse_method: str = "raw",
        hybrid: Optional[str] = None,
        image_output: Optional[str] = None,
        sanitize: Optional[bool] = None,
    ):
        self.outlines = extract_pdf_outlines(binary if binary is not None else filepath)

        if not self.api_url:
            raise RuntimeError("[OpenDataLoader] OPENDATALOADER_APISERVER is not configured. Please start the opendataloader service and set the env var.")

        # Render page images locally — used by _make_line_tag() and crop().
        # The image rendering stays on the RAGFlow host; only the Java conversion
        # runs inside the opendataloader service container.
        try:
            if binary is not None:
                src = BytesIO(binary) if isinstance(binary, (bytes, bytearray)) else binary
                self.__images__(src, zoomin=1)
            else:
                self.__images__(str(filepath), zoomin=1)
        except Exception as e:
            self.logger.warning(f"[OpenDataLoader] render pages failed: {e}")

        # Build page-height lookup for Docling TOPLEFT → bottom-left
        # coordinate conversion.  Pages are 1-based in Docling.
        page_heights: dict[int, float] = {}
        for i, img in enumerate(self.page_images):
            if img is not None:
                page_heights[i + 1] = float(img.size[1])

        # Read PDF bytes for the multipart upload
        if binary is not None:
            pdf_bytes = binary if isinstance(binary, (bytes, bytearray)) else binary.getvalue()
        else:
            with open(filepath, "rb") as fh:
                pdf_bytes = fh.read()

        filename = Path(str(filepath)).name or "input.pdf"

        if callback:
            callback(0.1, f"[OpenDataLoader] Sending '{filename}' to service")

        # Docling Fast Server endpoint: POST /v1/convert/file
        # - files: PDF file (multipart/form-data)
        # - page_ranges: optional, e.g. "1-5"
        form_data: dict[str, str] = {"to_formats": '["json", "md"]'}
        if hybrid:
            form_data["hybrid"] = hybrid
        if image_output:
            form_data["image_output"] = image_output
        if sanitize is not None:
            form_data["sanitize"] = "true" if sanitize else "false"

        headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
        last_exc: Exception | None = None
        for attempt in range(1, 4):
            try:
                self.logger.info(f"[OpenDataLoader] POST {self.api_url}/v1/convert/file for '{filename}' (attempt {attempt})")
                resp = requests.post(
                    url=f"{self.api_url}/v1/convert/file",
                    files={"files": (filename, pdf_bytes, "application/pdf")},
                    data=form_data,
                    headers=headers,
                    timeout=self.timeout,
                )
                resp.raise_for_status()
                result = resp.json()
                break
            except Exception as exc:
                last_exc = exc
                self.logger.warning(f"[OpenDataLoader] attempt {attempt} failed: {exc}")
        else:
            raise RuntimeError(f"[OpenDataLoader] service call failed after 3 attempts: {last_exc}") from last_exc

        if callback:
            callback(0.7, "[OpenDataLoader] Processing response")

        # /v1/convert/file response structure (nested):
        #   {"status": ..., "document": {"json_content": {DoclingDocument}}, "processing_time": ...}
        json_doc = None
        md_text = None
        if isinstance(result, dict):
            doc = result.get("document")
            inner = None
            if isinstance(doc, dict):
                inner = doc.get("json_content")
                if isinstance(inner, str):
                    try:
                        inner = json.loads(inner)
                    except (json.JSONDecodeError, TypeError):
                        self.logger.warning("[OpenDataLoader] json_content is a string but not valid JSON")
                        inner = None

            # Use DoclingDocument if found at inner, otherwise fall back
            if isinstance(inner, dict) and ("texts" in inner or "tables" in inner or "pictures" in inner):
                json_doc = _normalize_docling_response(inner, page_heights)
            elif isinstance(doc, dict) and ("texts" in doc or "tables" in doc or "pictures" in doc):
                json_doc = _normalize_docling_response(doc, page_heights)
            elif isinstance(result, dict) and ("texts" in result or "tables" in result or "pictures" in result):
                json_doc = _normalize_docling_response(result, page_heights)
            else:
                # Legacy / flat format
                json_doc = result.get("json_doc")
                md_text = result.get("md_text")
        elif isinstance(result, str):
            md_text = result

        sections: list[tuple[str, ...]] = []
        tables: list = []
        if json_doc is not None:
            sections, tables = self._transfer_from_json(json_doc, parse_method=parse_method)
            self.logger.info(f"[OpenDataLoader] Extracted {len(sections)} sections, {len(tables)} tables")
        if not sections and md_text:
            sections = self._sections_from_markdown(md_text, parse_method=parse_method)

        if callback:
            callback(1.0, f"[OpenDataLoader] Done. Sections: {len(sections)}, Tables: {len(tables)}")

        return sections, tables


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    parser = OpenDataLoaderParser()
    print("OpenDataLoader service reachable:", parser.check_installation())