X-Financial/server/src/app/services/attachment_association_jobs.py

from __future__ import annotations

import hashlib
from collections.abc import Callable
from contextlib import contextmanager
from threading import Lock
from typing import Any

from sqlalchemy import text
from sqlalchemy.orm import Session, sessionmaker

from app.api.deps import CurrentUserContext
from app.schemas.attachment_association_job import (
    AttachmentAssociationJobCreate,
    AttachmentAssociationJobRead,
)
from app.schemas.receipt_folder import ReceiptFolderDetailRead
from app.services.attachment_association_job_store import (
    ClaimedAttachmentAssociationJob,
    claim_persistent_job,
    create_persistent_job,
    get_authorized_persistent_job,
    job_to_read,
    update_persistent_job,
)
from app.services.expense_receipt_association import ExpenseReceiptAssociationService
from app.services.receipt_folder import ReceiptFolderService

SessionFactory = sessionmaker[Session] | Callable[[], Session]

_receipt_locks: dict[str, Lock] = {}
_job_claim_locks: dict[str, Lock] = {}
_receipt_locks_guard = Lock()


def clear_attachment_association_jobs_for_tests() -> None:
    """测试数据库按用例隔离；这里只清理进程级互斥锁。"""
    with _receipt_locks_guard:
        _receipt_locks.clear()
        _job_claim_locks.clear()


def create_attachment_association_job(
    payload: AttachmentAssociationJobCreate,
    current_user: CurrentUserContext,
    db: Session,
) -> AttachmentAssociationJobRead:
    return create_persistent_job(db, payload, current_user)


def get_attachment_association_job(
    job_id: str,
    current_user: CurrentUserContext,
    db: Session,
) -> AttachmentAssociationJobRead | None:
    job = get_authorized_persistent_job(db, job_id, current_user)
    return job_to_read(job) if job is not None else None


def run_attachment_association_job(
    job_id: str,
    session_factory: SessionFactory,
) -> None:
    with _receipt_locks_guard:
        claim_lock = _job_claim_locks.setdefault(str(job_id or "").strip(), Lock())
    with claim_lock:
        with session_factory() as claim_db:
            claimed = claim_persistent_job(claim_db, job_id)
    if claimed is None:
        return

    try:
        with session_factory() as db:
            with _receipt_execution_locks(db, claimed):
                result = AttachmentAssociationJobRunner(db).run(
                    receipt_ids=claimed.receipt_ids,
                    current_user=claimed.current_user,
                )
        requires_confirmation = bool(result.get("requires_confirmation"))
        if requires_confirmation:
            exceptions = [
                str(item) for item in list(result.get("exceptions") or []) if str(item).strip()
            ]
            message = exceptions[0] if exceptions else "匹配结果需要确认，系统未修改任何报销数据。"
        else:
            uploaded_count = int(result.get("uploaded_count") or 0)
            skipped_count = int(result.get("skipped_count") or 0)
            if uploaded_count == 0 and skipped_count > 0:
                message = "票据此前已经归集到目标草稿，本次未重复写入。"
            else:
                message = (
                    f"已自动关联到 {result.get('claim_no') or '报销草稿'}，"
                    f"成功归集 {uploaded_count} 份附件。"
                )
        with session_factory() as update_db:
            update_persistent_job(
                update_db,
                job_id,
                expected_attempt_count=claimed.attempt_count,
                status="succeeded",
                message=message,
                error="",
                **_job_result_updates(result),
            )
    except Exception as exc:
        message = str(exc).strip() or "自动关联任务执行失败，请稍后重试。"
        with session_factory() as update_db:
            update_persistent_job(
                update_db,
                job_id,
                expected_attempt_count=claimed.attempt_count,
                status="failed",
                message=message,
                error=message,
                resolution="failed",
            )


class AttachmentAssociationJobRunner:
    def __init__(self, db: Session) -> None:
        self.db = db
        self.receipt_service = ReceiptFolderService()

    def run(
        self,
        *,
        receipt_ids: list[str],
        current_user: CurrentUserContext,
    ) -> dict[str, Any]:
        receipts = self._load_receipts(receipt_ids, current_user)
        return ExpenseReceiptAssociationService(self.db).associate(
            receipts=receipts,
            current_user=current_user,
        )

    def _load_receipts(
        self,
        receipt_ids: list[str],
        current_user: CurrentUserContext,
    ) -> list[ReceiptFolderDetailRead]:
        receipts: list[ReceiptFolderDetailRead] = []
        normalized_ids = list(
            dict.fromkeys(
                str(item or "").strip() for item in receipt_ids if str(item or "").strip()
            )
        )
        for receipt_id in normalized_ids:
            try:
                receipts.append(self.receipt_service.get_receipt(receipt_id, current_user))
            except FileNotFoundError as exc:
                raise ValueError("当前附件没有持久化票据记录，请重新上传后再试。") from exc
        if not receipts:
            raise ValueError("当前附件没有持久化票据记录，请重新上传后再试。")
        return receipts


@contextmanager
def _receipt_execution_locks(
    db: Session,
    claimed: ClaimedAttachmentAssociationJob,
):
    lock_keys = sorted(
        {
            f"{claimed.current_user.tenant_id}:{receipt_id}"
            for receipt_id in claimed.receipt_ids
            if receipt_id
        }
    )
    with _receipt_locks_guard:
        process_locks = [_receipt_locks.setdefault(key, Lock()) for key in lock_keys]
    for process_lock in process_locks:
        process_lock.acquire()

    advisory_keys: list[int] = []
    try:
        if db.get_bind().dialect.name == "postgresql":
            advisory_keys = [_advisory_lock_key(key) for key in lock_keys]
            for advisory_key in advisory_keys:
                db.execute(
                    text("SELECT pg_advisory_lock(:lock_key)"),
                    {"lock_key": advisory_key},
                )
        yield
    finally:
        try:
            if db.get_bind().dialect.name == "postgresql":
                # 业务异常可能让当前事务失效；会话级 advisory lock 必须在可用事务中解锁。
                db.rollback()
                for advisory_key in reversed(advisory_keys):
                    db.execute(
                        text("SELECT pg_advisory_unlock(:lock_key)"),
                        {"lock_key": advisory_key},
                    )
        except Exception:
            # 连接关闭也会释放会话级锁，不能让解锁异常覆盖原始业务异常。
            db.rollback()
        finally:
            for process_lock in reversed(process_locks):
                process_lock.release()


def _advisory_lock_key(value: str) -> int:
    return int.from_bytes(
        hashlib.sha256(value.encode("utf-8")).digest()[:8],
        byteorder="big",
        signed=True,
    )


def _job_result_updates(result: dict[str, Any]) -> dict[str, Any]:
    return {
        "claim_id": str(result.get("claim_id") or ""),
        "claim_no": str(result.get("claim_no") or ""),
        "uploaded_count": int(result.get("uploaded_count") or 0),
        "skipped_count": int(result.get("skipped_count") or 0),
        "resolution": str(result.get("resolution") or ""),
        "requires_confirmation": bool(result.get("requires_confirmation")),
        "expense_case_id": str(result.get("expense_case_id") or ""),
        "application_claim_id": str(result.get("application_claim_id") or ""),
        "application_claim_no": str(result.get("application_claim_no") or ""),
        "confidence": str(result.get("confidence") or ""),
        "confidence_score": float(result.get("confidence_score") or 0.0),
        "match_reasons_json": [
            str(item) for item in list(result.get("match_reasons") or [])
        ],
        "exceptions_json": [str(item) for item in list(result.get("exceptions") or [])],
        "missing_fields_json": [
            str(item) for item in list(result.get("missing_fields") or [])
        ],
        "risk_items_json": [str(item) for item in list(result.get("risk_items") or [])],
        "candidates_json": [
            dict(item) for item in list(result.get("candidates") or []) if isinstance(item, dict)
        ],
        "draft_payload_json": (
            dict(result["draft_payload"])
            if isinstance(result.get("draft_payload"), dict)
            else None
        ),
    }