app/rag/academic_parser.py

"""Layout-aware academic PDF parsing using PyMuPDF4LLM.

The parser preserves page, reading order, section hierarchy, element type, and
normalized coordinates. It deliberately performs no embedding or storage so
the parse can be inspected before any derived index is replaced.
"""

from __future__ import annotations

import re
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Iterable

import pymupdf
import pymupdf4llm

from app.rag.models import (
    AcademicDocument,
    BoundingBox,
    EvidenceRelation,
    EvidenceType,
    EvidenceUnit,
    ParsedAcademicDocument,
    RelationType,
    SourceKind,
    make_evidence_id,
    normalize_evidence_text,
)


_BOX_TYPES: dict[str, EvidenceType] = {
    "title": EvidenceType.TITLE,
    "section-header": EvidenceType.HEADING,
    "text": EvidenceType.PARAGRAPH,
    "list-item": EvidenceType.LIST,
    "table": EvidenceType.TABLE,
    "picture": EvidenceType.FIGURE,
    "image": EvidenceType.FIGURE,
    "figure": EvidenceType.FIGURE,
    "plot": EvidenceType.PLOT,
    "diagram": EvidenceType.DIAGRAM,
    "formula": EvidenceType.FORMULA,
    "caption": EvidenceType.CAPTION,
    "footnote": EvidenceType.FOOTNOTE,
}
_IGNORED_BOX_TYPES = {"page-header", "page-footer"}
_REFERENCE_HEADING = re.compile(r"^(references|bibliography|works cited)\b", re.IGNORECASE)
_ABSTRACT_HEADING = re.compile(r"^abstract\b", re.IGNORECASE)
_MARKDOWN_HEADING = re.compile(r"^(#{1,6})\s+(.*)$")
_MARKDOWN_IMAGE = re.compile(r"!\[[^]]*]\([^)]*\)")


@dataclass(frozen=True)
class _PageInfo:
    number: int
    width: float
    height: float


def is_reference_section_heading(value: str) -> bool:
    """Recognize bibliography headings even when the PDF converter kept emphasis."""

    plain = re.sub(r"[*_`]+", "", value or "").strip()
    return bool(_REFERENCE_HEADING.match(plain))


class AcademicDocumentParser:
    """Convert a PDF into deterministic evidence units and relations."""

    def parse_pdf(
        self,
        file_bytes: bytes,
        filename: str,
        project_id: str,
        document_id: str,
    ) -> ParsedAcademicDocument:
        doc = pymupdf.open(stream=file_bytes, filetype="pdf")
        try:
            pages = pymupdf4llm.to_markdown(
                doc,
                page_chunks=True,
                extract_words=True,
                header=False,
                footer=False,
                show_progress=False,
            )
            if not isinstance(pages, list):
                raise RuntimeError("PyMuPDF4LLM did not return page chunks")
            page_info = {
                index + 1: _PageInfo(index + 1, float(page.rect.width), float(page.rect.height))
                for index, page in enumerate(doc)
            }
            return self._materialize(
                pages=pages,
                page_info=page_info,
                metadata=dict(doc.metadata or {}),
                filename=filename,
                project_id=project_id,
                document_id=document_id,
            )
        finally:
            doc.close()

    def parse_text(
        self,
        text: str,
        filename: str,
        project_id: str,
        document_id: str,
    ) -> ParsedAcademicDocument:
        """Structure non-PDF text without reintroducing fixed-width chunks."""

        units: list[EvidenceUnit] = []
        section_path: list[str] = []
        title = filename.rsplit(".", 1)[0]
        ordinal = 0
        for block in (part.strip() for part in re.split(r"\n\s*\n+", text or "")):
            if not block:
                continue
            level, heading_text = _heading(block)
            is_heading = bool(_MARKDOWN_HEADING.match(block))
            if is_heading:
                section_path = _update_section_path(section_path, level, heading_text)
                element_type = EvidenceType.HEADING
                content = heading_text
                if not units and heading_text:
                    title = heading_text
            else:
                element_type = EvidenceType.PARAGRAPH
                content = _strip_markdown(block)
            for fragment in _semantic_fragments(content, element_type):
                evidence_id = make_evidence_id(project_id, document_id, 1, None, element_type, fragment)
                units.append(
                    EvidenceUnit(
                        evidence_id=evidence_id,
                        project_id=project_id,
                        document_id=document_id,
                        element_type=element_type,
                        page_start=1,
                        page_end=1,
                        section_path=list(section_path),
                        ordinal=ordinal,
                        raw_text=fragment,
                        retrieval_text=_retrieval_text(
                            title=title,
                            section_path=section_path,
                            element_type=element_type,
                            content=fragment,
                        ),
                        quality_flags=["page_geometry_unavailable"],
                    )
                )
                ordinal += 1
        flags = ["page_geometry_unavailable"]
        if not units:
            flags.append("document_empty")
        relations = _build_relations(units, flags)
        cards, card_relations = _build_section_cards(
            units, project_id, document_id, title, start_ordinal=len(units)
        )
        units.extend(cards)
        relations.extend(card_relations)
        windows, window_relations = _build_local_windows(
            units, project_id, document_id, title, start_ordinal=len(units)
        )
        units.extend(windows)
        relations.extend(window_relations)
        return ParsedAcademicDocument(
            document=AcademicDocument(
                project_id=project_id,
                document_id=document_id,
                filename=filename,
                title=title,
                page_count=1 if units else 0,
                parse_quality="degraded" if units else "empty",
                quality_flags=flags,
            ),
            units=units,
            relations=_dedupe_relations(relations),
        )

    def _materialize(
        self,
        *,
        pages: list[dict[str, Any]],
        page_info: dict[int, _PageInfo],
        metadata: dict[str, Any],
        filename: str,
        project_id: str,
        document_id: str,
    ) -> ParsedAcademicDocument:
        quality_flags: list[str] = []
        units: list[EvidenceUnit] = []
        active_sections: list[str] = []
        first_title = ""
        ordinal = 0

        for page_index, page_chunk in enumerate(pages):
            metadata_page = page_chunk.get("metadata") or {}
            page_number = _product_page_number(metadata_page.get("page_number"), page_index + 1)
            info = page_info.get(page_number) or page_info.get(page_index + 1)
            if info is None:
                quality_flags.append(f"page_dimensions_missing:{page_number}")
                continue
            page_text = str(page_chunk.get("text") or "")
            boxes = sorted(page_chunk.get("page_boxes") or [], key=lambda item: int(item.get("index", 0)))
            if not page_text.strip():
                quality_flags.append(f"page_empty:{page_number}")
            if not boxes and page_text.strip():
                quality_flags.append(f"page_boxes_missing:{page_number}")
                boxes = [{"class": "text", "bbox": [0, 0, info.width, info.height], "pos": [0, len(page_text)], "index": 0}]

            for box in boxes:
                box_class = str(box.get("class") or "text").strip().lower()
                if box_class in _IGNORED_BOX_TYPES:
                    continue
                element_type = _BOX_TYPES.get(box_class, EvidenceType.PARAGRAPH)
                bbox = _normalized_box(box.get("bbox"), info)
                if bbox is None:
                    quality_flags.append(f"bbox_invalid:{page_number}:{box.get('index', 0)}")
                content = _box_text(page_text, box)

                if element_type is EvidenceType.HEADING:
                    heading_level, heading_text = _heading(content)
                    if not heading_text:
                        continue
                    active_sections = _update_section_path(active_sections, heading_level, heading_text)
                    content = heading_text
                elif element_type is EvidenceType.TITLE:
                    content = _strip_markdown(content)
                    if content and not first_title:
                        first_title = content

                fragments = _semantic_fragments(content, element_type)
                if not fragments and element_type in {
                    EvidenceType.FIGURE,
                    EvidenceType.PLOT,
                    EvidenceType.DIAGRAM,
                }:
                    fragments = [""]

                for fragment_index, fragment in enumerate(fragments):
                    effective_type = element_type
                    section_path = list(active_sections)
                    if section_path and is_reference_section_heading(section_path[-1]):
                        if element_type in {EvidenceType.PARAGRAPH, EvidenceType.LIST}:
                            effective_type = EvidenceType.REFERENCE
                    table_markdown = fragment if effective_type is EvidenceType.TABLE else ""
                    raw_text = "" if effective_type is EvidenceType.TABLE else fragment
                    content_key = fragment or f"{box_class}:{box.get('index', 0)}:{fragment_index}"
                    evidence_id = make_evidence_id(
                        project_id,
                        document_id,
                        page_number,
                        bbox,
                        effective_type,
                        content_key,
                    )
                    retrieval_text = _retrieval_text(
                        title=metadata.get("title") or first_title or filename,
                        section_path=section_path,
                        element_type=effective_type,
                        content=fragment,
                    )
                    unit_flags: list[str] = []
                    if not fragment:
                        unit_flags.append("visual_description_pending")
                    units.append(
                        EvidenceUnit(
                            evidence_id=evidence_id,
                            project_id=project_id,
                            document_id=document_id,
                            element_type=effective_type,
                            page_start=page_number,
                            page_end=page_number,
                            section_path=section_path,
                            ordinal=ordinal,
                            bbox_norm=bbox,
                            raw_text=raw_text,
                            retrieval_text=retrieval_text,
                            table_markdown=table_markdown,
                            quality_flags=unit_flags,
                            source_kind=SourceKind.PAPER,
                            metadata={
                                "box_class": box_class,
                                "box_index": int(box.get("index", 0)),
                                "fragment_index": fragment_index,
                            },
                        )
                    )
                    ordinal += 1

        title = normalize_evidence_text(str(metadata.get("title") or "")) or first_title or filename.rsplit(".", 1)[0]
        authors = _authors(str(metadata.get("author") or ""))
        abstract = _abstract_from_units(units)
        relations = _build_relations(units, quality_flags)
        section_cards, section_relations = _build_section_cards(
            units, project_id, document_id, title, start_ordinal=ordinal
        )
        units.extend(section_cards)
        relations.extend(section_relations)
        windows, window_relations = _build_local_windows(
            units, project_id, document_id, title, start_ordinal=len(units)
        )
        units.extend(windows)
        relations.extend(window_relations)

        leaf_units = [
            unit
            for unit in units
            if unit.element_type not in {EvidenceType.SECTION_CARD, EvidenceType.LOCAL_WINDOW}
        ]
        if not leaf_units:
            parse_quality: str = "empty"
            quality_flags.append("document_empty")
        elif any(
            flag.startswith(("page_empty:", "page_boxes_missing:", "page_dimensions_missing:"))
            for flag in quality_flags
        ):
            parse_quality = "degraded"
        else:
            parse_quality = "good"
        document = AcademicDocument(
            project_id=project_id,
            document_id=document_id,
            filename=filename,
            title=title,
            authors=authors,
            abstract=abstract,
            page_count=len(pages),
            parse_quality=parse_quality,
            quality_flags=sorted(set(quality_flags)),
        )
        return ParsedAcademicDocument(document=document, units=units, relations=_dedupe_relations(relations))


def _product_page_number(value: Any, fallback: int) -> int:
    """Use installed PyMuPDF4LLM's one-based page metadata, or index fallback."""

    try:
        number = int(value)
        return number if number >= 1 else fallback
    except (TypeError, ValueError):
        return fallback


def _normalized_box(value: Any, page: _PageInfo) -> BoundingBox | None:
    if not isinstance(value, (list, tuple)) or len(value) != 4 or page.width <= 0 or page.height <= 0:
        return None
    try:
        x0, y0, x1, y1 = (float(part) for part in value)
        x0 = min(max(x0 / page.width, 0.0), 1.0)
        y0 = min(max(y0 / page.height, 0.0), 1.0)
        x1 = min(max(x1 / page.width, 0.0), 1.0)
        y1 = min(max(y1 / page.height, 0.0), 1.0)
        if x1 <= x0 or y1 <= y0:
            return None
        return BoundingBox(x=x0, y=y0, w=x1 - x0, h=y1 - y0)
    except (TypeError, ValueError):
        return None


def _box_text(page_text: str, box: dict[str, Any]) -> str:
    pos = box.get("pos")
    if isinstance(pos, (list, tuple)) and len(pos) == 2:
        try:
            start = max(0, int(pos[0]))
            end = min(len(page_text), int(pos[1]))
            return page_text[start:end].strip()
        except (TypeError, ValueError):
            pass
    return ""


def _heading(content: str) -> tuple[int, str]:
    cleaned = content.strip()
    match = _MARKDOWN_HEADING.match(cleaned)
    if match:
        return len(match.group(1)), normalize_evidence_text(match.group(2))
    return 1, _strip_markdown(cleaned)


def _strip_markdown(content: str) -> str:
    content = _MARKDOWN_IMAGE.sub("", content or "")
    content = re.sub(r"^#{1,6}\s+", "", content.strip())
    # PyMuPDF4LLM wraps emphasized text (title/author lines are often bolded in
    # the source PDF) in markdown emphasis markers; leaving them in leaks
    # literal "**"/"`" characters into titles shown to the user.
    content = re.sub(r"[*_`]+", "", content)
    return normalize_evidence_text(content)


def _semantic_fragments(content: str, element_type: EvidenceType) -> list[str]:
    content = (content or "").strip()
    if not content:
        return []
    if element_type in {
        EvidenceType.TABLE,
        EvidenceType.FORMULA,
        EvidenceType.CAPTION,
        EvidenceType.TITLE,
        EvidenceType.HEADING,
        EvidenceType.FOOTNOTE,
        EvidenceType.REFERENCE,
    }:
        cleaned = content if element_type is EvidenceType.TABLE else _strip_markdown(content)
        return [cleaned] if cleaned else []

    parts = [part.strip() for part in re.split(r"\n\s*\n+", content) if part.strip()]
    fragments: list[str] = []
    for part in parts:
        cleaned = _strip_markdown(part)
        if not cleaned:
            continue
        # Layout boxes normally align with paragraphs. Exceptionally large boxes
        # are split on sentence boundaries without crossing their page/section.
        if len(cleaned) <= 1800:
            fragments.append(cleaned)
            continue
        sentences = re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])", cleaned)
        buffer: list[str] = []
        size = 0
        for sentence in sentences:
            if buffer and size + len(sentence) > 1400:
                fragments.append(" ".join(buffer))
                buffer, size = [], 0
            buffer.append(sentence)
            size += len(sentence) + 1
        if buffer:
            fragments.append(" ".join(buffer))
    return fragments


def _update_section_path(path: list[str], level: int, heading: str) -> list[str]:
    level = min(max(level, 1), 6)
    next_path = list(path[: level - 1])
    while len(next_path) < level - 1:
        next_path.append("Untitled section")
    next_path.append(heading)
    return next_path


def _retrieval_text(
    *,
    title: str,
    section_path: list[str],
    element_type: EvidenceType,
    content: str,
) -> str:
    context = [f"Paper: {normalize_evidence_text(title)}"]
    if section_path:
        context.append(f"Section: {' > '.join(section_path)}")
    context.append(f"Element: {element_type.value}")
    if content.strip():
        context.append(content.strip())
    return "\n".join(context)


def _authors(value: str) -> list[str]:
    if not value.strip():
        return []
    return [normalize_evidence_text(item) for item in re.split(r"[;,]", value) if item.strip()]


def _abstract_from_units(units: list[EvidenceUnit]) -> str:
    abstract_parts = [
        unit.raw_text
        for unit in units
        if unit.section_path and _ABSTRACT_HEADING.match(unit.section_path[-1]) and unit.raw_text
    ]
    return "\n".join(abstract_parts)[:6000]


def _build_relations(units: list[EvidenceUnit], quality_flags: list[str]) -> list[EvidenceRelation]:
    relations: list[EvidenceRelation] = []
    leaves = [
        unit
        for unit in units
        if unit.element_type not in {EvidenceType.SECTION_CARD, EvidenceType.LOCAL_WINDOW}
    ]
    for previous, current in zip(leaves, leaves[1:]):
        relations.append(
            EvidenceRelation(
                source_evidence_id=previous.evidence_id,
                target_evidence_id=current.evidence_id,
                relation_type=RelationType.NEXT,
                derivation="parser",
            )
        )
        relations.append(
            EvidenceRelation(
                source_evidence_id=current.evidence_id,
                target_evidence_id=previous.evidence_id,
                relation_type=RelationType.PREVIOUS,
                derivation="parser",
            )
        )

    by_page: dict[int, list[EvidenceUnit]] = defaultdict(list)
    for unit in leaves:
        by_page[unit.page_start].append(unit)
    visual_types = {EvidenceType.FIGURE, EvidenceType.PLOT, EvidenceType.DIAGRAM, EvidenceType.TABLE}
    for page_number, page_units in by_page.items():
        visuals = [unit for unit in page_units if unit.element_type in visual_types and unit.bbox_norm]
        captions = [unit for unit in page_units if unit.element_type is EvidenceType.CAPTION and unit.bbox_norm]
        for caption in captions:
            candidates = sorted(
                visuals,
                key=lambda visual: _vertical_distance(caption.bbox_norm, visual.bbox_norm),
            )
            if not candidates:
                quality_flags.append(f"caption_unlinked:{page_number}:{caption.evidence_id}")
                continue
            target = candidates[0]
            if _vertical_distance(caption.bbox_norm, target.bbox_norm) > 0.25:
                quality_flags.append(f"caption_unlinked:{page_number}:{caption.evidence_id}")
                continue
            target.caption = caption.raw_text
            target.retrieval_text = f"{target.retrieval_text}\nCaption: {caption.raw_text}".strip()
            relations.append(
                EvidenceRelation(
                    source_evidence_id=caption.evidence_id,
                    target_evidence_id=target.evidence_id,
                    relation_type=RelationType.CAPTION_OF,
                    confidence=0.9,
                    derivation="deterministic_linker",
                )
            )
    return relations


def _vertical_distance(first: BoundingBox | None, second: BoundingBox | None) -> float:
    if first is None or second is None:
        return 1.0
    first_mid = first.y + first.h / 2
    second_mid = second.y + second.h / 2
    return abs(first_mid - second_mid)


def _build_section_cards(
    units: list[EvidenceUnit],
    project_id: str,
    document_id: str,
    title: str,
    *,
    start_ordinal: int,
) -> tuple[list[EvidenceUnit], list[EvidenceRelation]]:
    grouped: dict[tuple[str, ...], list[EvidenceUnit]] = defaultdict(list)
    for unit in units:
        if unit.section_path and unit.element_type not in {EvidenceType.HEADING, EvidenceType.TITLE}:
            grouped[tuple(unit.section_path)].append(unit)
    cards: list[EvidenceUnit] = []
    relations: list[EvidenceRelation] = []
    for offset, (path, children) in enumerate(grouped.items()):
        content = " ".join(
            child.raw_text or child.table_markdown or child.caption or child.visual_description
            for child in children
            if child.index_text.strip()
        )
        preview = normalize_evidence_text(content)[:2400]
        first = children[0]
        evidence_id = make_evidence_id(
            project_id,
            document_id,
            first.page_start,
            None,
            EvidenceType.SECTION_CARD,
            " > ".join(path),
        )
        card = EvidenceUnit(
            evidence_id=evidence_id,
            project_id=project_id,
            document_id=document_id,
            element_type=EvidenceType.SECTION_CARD,
            page_start=min(child.page_start for child in children),
            page_end=max(child.page_end for child in children),
            section_path=list(path),
            ordinal=start_ordinal + offset,
            raw_text=preview,
            retrieval_text=_retrieval_text(
                title=title,
                section_path=list(path),
                element_type=EvidenceType.SECTION_CARD,
                content=preview,
            ),
            source_kind=SourceKind.DERIVED_SUMMARY,
            metadata={"child_count": len(children)},
        )
        cards.append(card)
        for child in children:
            child.parent_id = evidence_id
            relations.append(
                EvidenceRelation(
                    source_evidence_id=child.evidence_id,
                    target_evidence_id=evidence_id,
                    relation_type=RelationType.PARENT,
                    derivation="deterministic_linker",
                )
            )
    return cards, relations


def _build_local_windows(
    units: list[EvidenceUnit],
    project_id: str,
    document_id: str,
    title: str,
    *,
    start_ordinal: int,
) -> tuple[list[EvidenceUnit], list[EvidenceRelation]]:
    grouped: dict[tuple[str, ...], list[EvidenceUnit]] = defaultdict(list)
    eligible = {
        EvidenceType.PARAGRAPH,
        EvidenceType.LIST,
        EvidenceType.TABLE,
        EvidenceType.FORMULA,
        EvidenceType.CAPTION,
        EvidenceType.REFERENCE,
    }
    for unit in units:
        if unit.element_type in eligible and unit.section_path and unit.index_text.strip():
            grouped[tuple(unit.section_path)].append(unit)
    windows: list[EvidenceUnit] = []
    relations: list[EvidenceRelation] = []
    ordinal = start_ordinal
    for path, children in grouped.items():
        children.sort(key=lambda unit: unit.ordinal)
        for start in range(0, len(children), 3):
            group = children[start : start + 3]
            if len(group) < 2:
                continue
            content = "\n\n".join(child.raw_text or child.table_markdown for child in group)
            evidence_id = make_evidence_id(
                project_id,
                document_id,
                group[0].page_start,
                None,
                EvidenceType.LOCAL_WINDOW,
                "|".join(child.evidence_id for child in group),
            )
            window = EvidenceUnit(
                evidence_id=evidence_id,
                project_id=project_id,
                document_id=document_id,
                element_type=EvidenceType.LOCAL_WINDOW,
                page_start=min(child.page_start for child in group),
                page_end=max(child.page_end for child in group),
                parent_id=group[0].parent_id,
                section_path=list(path),
                ordinal=ordinal,
                raw_text=content,
                retrieval_text=_retrieval_text(
                    title=title,
                    section_path=list(path),
                    element_type=EvidenceType.LOCAL_WINDOW,
                    content=content,
                ),
                source_kind=SourceKind.DERIVED_SUMMARY,
                metadata={"child_evidence_ids": [child.evidence_id for child in group]},
            )
            windows.append(window)
            ordinal += 1
            for child in group:
                relations.append(
                    EvidenceRelation(
                        source_evidence_id=child.evidence_id,
                        target_evidence_id=window.evidence_id,
                        relation_type=RelationType.SUPPORTS,
                        derivation="deterministic_linker",
                    )
                )
    return windows, relations


def _dedupe_relations(relations: Iterable[EvidenceRelation]) -> list[EvidenceRelation]:
    seen: set[tuple[str, str, str]] = set()
    out: list[EvidenceRelation] = []
    for relation in relations:
        key = (
            relation.source_evidence_id,
            relation.target_evidence_id,
            str(relation.relation_type),
        )
        if key not in seen:
            seen.add(key)
            out.append(relation)
    return out