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

from __future__ import annotations

import hashlib
import io
import json
import re
import zipfile
from typing import Any, Callable, Mapping, Optional, Sequence
from urllib.parse import quote, unquote, urlparse

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

from app.utils.excel_export import DictLabelResolver, ExportColumn, ExportValidationError

ARCHIVE_ZIP_ATTACH_COL_MAX = 20
ARCHIVE_ZIP_MAX_FILES = 2000
ARCHIVE_ZIP_MAX_BYTES = 500 * 1024 * 1024


def sanitize_zip_segment(raw: str) -> str:
    value = re.sub(r'[\\/:*?"<>|]+', "_", str(raw or "").strip())
    value = re.sub(r"\s+", " ", value).strip()
    return value or "未命名"


def short_hash(text: str, *, n: int = 6) -> str:
    return hashlib.sha1(str(text or "").encode("utf-8")).hexdigest()[:n]


def guess_ext_from_url(url: str) -> str:
    path = unquote(urlparse(str(url or "").split("?", 1)[0]).path)
    name = path.rsplit("/", 1)[-1]
    if "." not in name:
        return "bin"
    ext = name.rsplit(".", 1)[-1].lower()
    if not ext or len(ext) > 8 or "/" in ext:
        return "bin"
    return ext


def build_attachment_filename(row_no: int, label: str, url: str) -> str:
    return (
        f"{int(row_no):02d}_{sanitize_zip_segment(label)}_"
        f"{short_hash(url)}.{guess_ext_from_url(url)}"
    )


def first_row_no_by_plan_id(rows: Sequence[Mapping[str, Any]]) -> dict[str, int]:
    out: dict[str, int] = {}
    for row in rows:
        plan_id = str(row.get("planId") or "").strip()
        if not plan_id or plan_id in out:
            continue
        try:
            out[plan_id] = int(row.get("rowNo") or 0)
        except (TypeError, ValueError):
            out[plan_id] = 0
    return out


def build_hyperlink_formula(rel_path: str, display: str) -> str:
    path = str(rel_path).replace('"', '""')
    text = sanitize_zip_segment(display).replace('"', '""')
    return f'=HYPERLINK("{path}","{text}")'


def parse_url_list(raw: Optional[str]) -> list[str]:
    if not raw:
        return []
    text = str(raw).strip()
    if not text:
        return []
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return [str(x).strip() for x in parsed if str(x).strip()]
    except json.JSONDecodeError:
        pass
    return [p.strip() for p in text.split(",") if p.strip()]


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 _default_download(url: str) -> bytes:
    from app.services.bx.storage import download_file_bytes

    return download_file_bytes(url)


def _attach_field(index: int) -> str:
    return f"attach{index}"


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)。"""
    download = download_fn or _default_download
    row_no_by_plan = first_row_no_by_plan_id(rows)
    root = f"计划归档({stamp})"
    attach_dir = f"{root}/附件"

    # plan_id -> list of (rel_path_in_zip_under_附件, label, formula_rel)
    plan_file_meta: dict[str, list[tuple[str, str, str]]] = {}
    zip_file_entries: list[tuple[str, bytes]] = []
    failures: list[str] = []
    total_bytes = 0
    file_count = 0

    for plan_id, items in (attachments_by_plan or {}).items():
        seq = row_no_by_plan.get(plan_id) or 0
        if seq <= 0:
            continue
        meta: list[tuple[str, str, str]] = []
        for item in items or []:
            url = str(item.get("url") or "").strip()
            if not url:
                continue
            label = str(item.get("label") or "附件")
            filename = build_attachment_filename(seq, label, url)
            zip_path = f"{attach_dir}/{filename}"
            try:
                content = download(url)
            except Exception as exc:  # noqa: BLE001 — 单文件失败不阻断整包
                failures.append(f"{label}: {url} ({exc})")
                continue
            file_count += 1
            total_bytes += len(content)
            if file_count > ARCHIVE_ZIP_MAX_FILES:
                raise ExportValidationError(
                    f"附件文件数超过 {ARCHIVE_ZIP_MAX_FILES}，请缩小筛选范围"
                )
            if total_bytes > ARCHIVE_ZIP_MAX_BYTES:
                raise ExportValidationError(
                    "附件总体积过大，请缩小筛选范围后重试"
                )
            zip_file_entries.append((zip_path, content))
            formula_rel = f"附件/{filename}"
            meta.append((zip_path, label, formula_rel))
        if meta:
            plan_file_meta[plan_id] = meta

    max_att = 0
    for meta in plan_file_meta.values():
        max_att = max(max_att, len(meta))
    k = min(ARCHIVE_ZIP_ATTACH_COL_MAX, max_att)

    attach_columns = [
        ExportColumn(title=f"附件_{i}", field=_attach_field(i)) for i in range(1, k + 1)
    ]
    all_columns = list(columns) + attach_columns

    enriched_rows: list[dict[str, Any]] = []
    seen_first: set[str] = set()
    for row in rows:
        data = dict(row)
        plan_id = str(data.get("planId") or "").strip()
        for i in range(1, k + 1):
            data[_attach_field(i)] = ""
        if plan_id and plan_id not in seen_first:
            seen_first.add(plan_id)
            meta = plan_file_meta.get(plan_id) or []
            for i, (_zip_path, label, formula_rel) in enumerate(meta[:k], start=1):
                data[_attach_field(i)] = build_hyperlink_formula(formula_rel, label)
        enriched_rows.append(data)

    if not enriched_rows:
        raise ExportValidationError("没有可导出的数据")
    if not all_columns:
        raise ExportValidationError("没有可导出的列")

    resolver = DictLabelResolver(db, all_columns, tenant_no=tenant_no)
    from openpyxl import Workbook

    wb = Workbook()
    ws = wb.active
    ws.title = "计划归档"
    ws.append([c.title for c in all_columns])
    for row in enriched_rows:
        cells = resolver.row_to_cells(row, all_columns)
        ws.append(cells)

    xlsx_buf = io.BytesIO()
    wb.save(xlsx_buf)
    xlsx_bytes = xlsx_buf.getvalue()

    out = io.BytesIO()
    with zipfile.ZipFile(out, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr(f"{root}/(总表)计划归档.xlsx", xlsx_bytes)
        for path, content in zip_file_entries:
            zf.writestr(path, content)
        if failures:
            note = "以下附件下载失败，已跳过：\n" + "\n".join(failures)
            zf.writestr(f"{root}/导出说明.txt", note.encode("utf-8"))
    return out.getvalue(), failures


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:
    data, _failures = build_archive_zip_bytes(
        db=db,
        tenant_no=tenant_no,
        columns=columns,
        rows=rows,
        attachments_by_plan=attachments_by_plan,
        stamp=stamp,
        download_fn=download_fn,
    )
    filename = f"计划归档({stamp}).zip"
    encoded = quote(filename)
    disposition = f"attachment; filename={encoded}; filename*=UTF-8''{encoded}"
    return StreamingResponse(
        io.BytesIO(data),
        media_type="application/zip",
        headers={"Content-Disposition": disposition},
    )