shanxia2.0/backend/app/services/pl_plan_archive_zip_export.py

"""运输计划归档 ZIP 导出（计划特有附件收集 + 公共打包）。"""

from __future__ import annotations

from typing import Any, Callable, Mapping, Optional, Sequence

from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.services.export_attachment_zip import (
    ARCHIVE_ZIP_ATTACH_COL_MAX,
    ARCHIVE_ZIP_MAX_BYTES,
    ARCHIVE_ZIP_MAX_FILES,
    build_attachment_filename,
    build_attachment_zip_bytes,
    build_attachment_zip_streaming_response,
    build_hyperlink_formula,
    first_row_no_by_entity_id,
    parse_url_list,
    sanitize_zip_segment,
)
from app.utils.excel_export import ExportColumn

__all__ = [
    "ARCHIVE_ZIP_ATTACH_COL_MAX",
    "ARCHIVE_ZIP_MAX_BYTES",
    "ARCHIVE_ZIP_MAX_FILES",
    "sanitize_zip_segment",
    "build_attachment_filename",
    "first_row_no_by_plan_id",
    "build_hyperlink_formula",
    "parse_url_list",
    "collect_hazard_cert_items",
    "merge_plan_attachment_items",
    "load_plan_products_by_plan_ids",
    "build_archive_zip_bytes",
    "build_archive_zip_streaming_response",
]


def first_row_no_by_plan_id(rows: Sequence[Mapping[str, Any]]) -> dict[str, int]:
    return first_row_no_by_entity_id(rows, id_field="planId")


def collect_hazard_cert_items(products: Sequence[Any]) -> list[dict[str, str]]:
    out: list[dict[str, str]] = []
    for product in products or []:
        pname = sanitize_zip_segment(getattr(product, "product_name", None) or "")
        specs = (
            ("购买证", getattr(product, "purchase_cert_url", None)),
            ("运输证", getattr(product, "transport_cert_url", None)),
            ("合法使用证明", getattr(product, "legal_use_cert_url", None)),
        )
        for base_label, raw in specs:
            for url in parse_url_list(raw):
                label = base_label if not pname or pname == "未命名" else f"{base_label}·{pname}"
                out.append({"url": url, "label": label})
    return out


def merge_plan_attachment_items(
    document_items: Sequence[Mapping[str, str]],
    hazard_items: Sequence[Mapping[str, str]],
) -> list[dict[str, str]]:
    seen: set[str] = set()
    out: list[dict[str, str]] = []
    for item in list(document_items or []) + list(hazard_items or []):
        url = str(item.get("url") or "").strip()
        if not url or url in seen:
            continue
        seen.add(url)
        out.append({"url": url, "label": str(item.get("label") or "附件")})
    return out


def load_plan_products_by_plan_ids(db, *, tenant_no: str, plan_ids: list[str]):
    from app.models.pl_models import PlPlanProduct

    if not plan_ids:
        return {}
    rows = (
        db.query(PlPlanProduct)
        .filter(
            PlPlanProduct.plan_id.in_(plan_ids),
            PlPlanProduct.tenant_no == tenant_no,
            PlPlanProduct.del_flag == 0,
        )
        .order_by(PlPlanProduct.order_num, PlPlanProduct.create_time)
        .all()
    )
    grouped: dict[str, list] = {pid: [] for pid in plan_ids}
    for row in rows:
        grouped.setdefault(row.plan_id, []).append(row)
    return grouped


def build_archive_zip_bytes(
    *,
    db: Optional[Session],
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    attachments_by_plan: Mapping[str, Sequence[Mapping[str, str]]],
    stamp: str,
    download_fn: Optional[Callable[[str], bytes]] = None,
) -> tuple[bytes, list[str]]:
    """打包归档 ZIP：xlsx + 附件/；返回 (zip_bytes, failure_messages)。"""
    return build_attachment_zip_bytes(
        db=db,
        tenant_no=tenant_no,
        columns=columns,
        rows=rows,
        attachments_by_id=attachments_by_plan,
        stamp=stamp,
        package_label="计划归档",
        id_field="planId",
        sheet_title="计划归档",
        xlsx_basename="(总表)计划归档.xlsx",
        download_fn=download_fn,
    )


def build_archive_zip_streaming_response(
    *,
    db: Session,
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    attachments_by_plan: Mapping[str, Sequence[Mapping[str, str]]],
    stamp: str,
    download_fn: Optional[Callable[[str], bytes]] = None,
) -> StreamingResponse:
    return build_attachment_zip_streaming_response(
        db=db,
        tenant_no=tenant_no,
        columns=columns,
        rows=rows,
        attachments_by_id=attachments_by_plan,
        stamp=stamp,
        package_label="计划归档",
        id_field="planId",
        sheet_title="计划归档",
        xlsx_basename="(总表)计划归档.xlsx",
        download_fn=download_fn,
    )