baodan/api/insurance/document/service.py

"""Document IR 摄取、持久化与授权查询。"""
from __future__ import annotations

import json
import os

from insurance.db.compat import db
from insurance.document.ingest import IngestResult, ingest_pdf
from insurance.models.insurance_document import InsuranceDocument, InsuranceDocumentPage, InsuranceDocumentTable


def ingest_and_persist(file_path: str, user_id: str, *, original_name: str = "", **ingest_options) -> InsuranceDocument:
    if "ocr_provider" not in ingest_options:
        from insurance.document.ocr import default_ocr_provider
        ingest_options["ocr_provider"] = default_ocr_provider()
    result = ingest_pdf(file_path, storage_key=os.path.abspath(file_path), **ingest_options)
    try:
        return persist_ingest_result(result, user_id, original_name=original_name or os.path.basename(file_path))
    except Exception:
        db.session.rollback()
        raise


def ingest_for_processing(file_path: str, user_id: str, *, original_name: str = "", **ingest_options):
    """摄取并持久化，同时把本次不可变 IR 返回给当前解析任务。"""
    if "ocr_provider" not in ingest_options:
        from insurance.document.ocr import default_ocr_provider
        ingest_options["ocr_provider"] = default_ocr_provider()
    result = ingest_pdf(file_path, storage_key=os.path.abspath(file_path), **ingest_options)
    try:
        document = persist_ingest_result(
            result,
            user_id,
            original_name=original_name or os.path.basename(file_path),
        )
    except Exception:
        db.session.rollback()
        raise
    return document, result


def persist_ingest_result(result: IngestResult, user_id: str, *, original_name: str) -> InsuranceDocument:
    ir = result.document_ir
    metadata = ir["document"]
    document = InsuranceDocument.query.filter_by(user_id=str(user_id), sha256=metadata["sha256"]).first()
    if document is None:
        document = InsuranceDocument(user_id=str(user_id), sha256=metadata["sha256"])
        db.session.add(document)
    document.storage_key = result.storage_key
    document.original_name = os.path.basename(str(original_name or "document.pdf"))[:255]
    document.mime_type = metadata["mimeType"]
    document.file_size = os.path.getsize(result.storage_key)
    document.page_count = metadata["pageCount"]
    document.document_type = "plan_illustration"
    document.pdf_kind = metadata["pdfKind"]
    document.status = metadata["status"]
    document.parser_version = metadata["parserVersion"]
    document.ocr_version = metadata.get("ocrVersion")
    document.profile_code = metadata.get("profileCode")
    document.profile_version = metadata.get("profileVersion")
    document.quality_json = json.dumps(ir["quality"], ensure_ascii=False)
    document.processing_log_json = json.dumps(result.processing_log, ensure_ascii=False)
    db.session.flush()

    InsuranceDocumentPage.query.filter_by(document_id=document.id).delete(synchronize_session=False)
    InsuranceDocumentTable.query.filter_by(document_id=document.id).delete(synchronize_session=False)
    for page in ir["pages"]:
        db.session.add(InsuranceDocumentPage(
            document_id=document.id,
            page_number=page["pageNumber"],
            width=page["width"],
            height=page["height"],
            native_text_quality=page.get("nativeTextQuality"),
            ocr_quality=page.get("ocrQuality"),
            page_class=page["pageClass"],
            text_blocks_json=json.dumps(page.get("textBlocks") or [], ensure_ascii=False),
            ocr_blocks_json=json.dumps(page.get("ocrBlocks") or [], ensure_ascii=False),
            images_json="[]",
        ))
    for table in ir["tables"]:
        db.session.add(InsuranceDocumentTable(
            document_id=document.id,
            table_id=table["tableId"],
            scenario_type=table.get("scenarioType") or "unknown",
            headers_json=json.dumps(table.get("headers") or [], ensure_ascii=False),
            rows_json=json.dumps(table.get("rows") or [], ensure_ascii=False),
            source_pages_json=json.dumps(table.get("sourcePages") or [], ensure_ascii=False),
            continuation_of=table.get("continuationOf"),
            bbox_json=json.dumps(table.get("bbox"), ensure_ascii=False),
            conflicts_json=json.dumps(table.get("conflicts") or [], ensure_ascii=False),
        ))
    db.session.commit()
    return document


def get_owned_document(document_id: int, user_id: str) -> InsuranceDocument | None:
    return InsuranceDocument.query.filter_by(id=document_id, user_id=str(user_id)).first()


def get_owned_page(document_id: int, page_number: int, user_id: str) -> tuple[InsuranceDocument | None, InsuranceDocumentPage | None]:
    document = get_owned_document(document_id, user_id)
    if not document:
        return None, None
    page = InsuranceDocumentPage.query.filter_by(document_id=document.id, page_number=page_number).first()
    return document, page


def load_document_ir(document: InsuranceDocument) -> dict:
    """从持久化记录恢复不可变 Document IR，供异步提取任务消费。"""
    pages = InsuranceDocumentPage.query.filter_by(document_id=document.id).order_by(
        InsuranceDocumentPage.page_number
    ).all()
    tables = InsuranceDocumentTable.query.filter_by(document_id=document.id).order_by(
        InsuranceDocumentTable.id
    ).all()
    page_payloads = []
    for item in pages:
        payload = item.to_public_dict()
        payload.pop("documentId", None)
        payload.pop("images", None)
        page_payloads.append(payload)
    table_payloads = []
    for item in tables:
        payload = item.to_public_dict()
        payload.pop("documentId", None)
        table_payloads.append(payload)
    try:
        quality = json.loads(document.quality_json or "{}")
    except (TypeError, ValueError):
        quality = {}
    quality.setdefault("nativePageCount", sum(1 for item in pages if item.page_class == "native"))
    quality.setdefault("ocrPageCount", sum(1 for item in pages if item.ocr_blocks_json not in (None, "", "[]")))
    quality.setdefault("lowQualityPages", [
        item.page_number for item in pages if item.page_class in {"scanned", "low_quality"}
    ])
    return {
        "schemaVersion": "document-ir-v1",
        "document": {
            "sha256": document.sha256,
            "mimeType": document.mime_type,
            "pageCount": document.page_count,
            "pdfKind": document.pdf_kind,
            "status": document.status,
            "parserVersion": document.parser_version,
            "ocrVersion": document.ocr_version,
            "profileCode": document.profile_code,
            "profileVersion": document.profile_version,
        },
        "pages": page_payloads,
        "tables": table_payloads,
        "quality": quality,
    }


def find_owned_evidence(document_id: int, user_id: str, field_path: str) -> tuple[InsuranceDocument | None, list[dict]]:
    document = get_owned_document(document_id, user_id)
    if not document:
        return None, []
    try:
        from insurance.models.plan_snapshot import InsuranceFieldEvidence, InsurancePlanSnapshot

        structured = db.session.query(InsuranceFieldEvidence).join(
            InsurancePlanSnapshot,
            InsurancePlanSnapshot.id == InsuranceFieldEvidence.snapshot_id,
        ).filter(
            InsurancePlanSnapshot.user_id == str(user_id),
            InsuranceFieldEvidence.document_id == document_id,
            InsuranceFieldEvidence.field_path == field_path,
        ).order_by(InsuranceFieldEvidence.id.desc()).limit(200).all()
        if structured:
            return document, [item.to_public_dict() for item in structured]
    except Exception:
        db.session.rollback()
    entries = []

    from insurance.models.ppt_session import PptSession
    sessions = PptSession.query.filter_by(user_id=str(user_id)).order_by(PptSession.updated_at.desc()).limit(100).all()
    for session in sessions:
        try:
            extractions = json.loads(session.extractions_json or "[]")
        except (TypeError, ValueError):
            continue
        for extraction in extractions:
            if _safe_int(extraction.get("documentId")) == document_id:
                entries.extend(extraction.get("evidence") or [])

    from insurance.models.poster_case_upload import PosterCaseUpload
    cases = PosterCaseUpload.query.filter_by(user_id=str(user_id)).order_by(PosterCaseUpload.created_at.desc()).limit(100).all()
    for case in cases:
        try:
            parsed = json.loads(case.parsed_data or "{}")
        except (TypeError, ValueError):
            continue
        meta = parsed.get("meta") or {}
        if _safe_int(meta.get("documentId")) == document_id:
            entries.extend(meta.get("evidence") or [])

    unique = {}
    for entry in entries:
        if not isinstance(entry, dict) or entry.get("fieldPath") != field_path:
            continue
        key = json.dumps(entry.get("evidence") or {}, sort_keys=True, ensure_ascii=False)
        unique[key] = entry
    return document, list(unique.values())


def _safe_int(value) -> int | None:
    try:
        return int(value)
    except (TypeError, ValueError):
        return None