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

"""列表导出：Excel 总表 + 附件目录 ZIP（计划归档 / 充值申请等共用）。"""

from __future__ import annotations

import hashlib
import io
import json
import re
import zipfile
from typing import Any, Callable, Iterator, Mapping, Optional, Sequence, Tuple
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

EXPORT_ZIP_ATTACH_COL_MAX = 20
EXPORT_ZIP_MAX_FILES = 2000
EXPORT_ZIP_MAX_BYTES = 500 * 1024 * 1024

# 兼容旧常量名
ARCHIVE_ZIP_ATTACH_COL_MAX = EXPORT_ZIP_ATTACH_COL_MAX
ARCHIVE_ZIP_MAX_FILES = EXPORT_ZIP_MAX_FILES
ARCHIVE_ZIP_MAX_BYTES = EXPORT_ZIP_MAX_BYTES


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_entity_id(
    rows: Sequence[Mapping[str, Any]], *, id_field: str
) -> dict[str, int]:
    out: dict[str, int] = {}
    for row in rows:
        entity_id = str(row.get(id_field) or "").strip()
        if not entity_id or entity_id in out:
            continue
        try:
            out[entity_id] = int(row.get("rowNo") or 0)
        except (TypeError, ValueError):
            out[entity_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 urls_to_attachment_items(
    raw: Optional[str], *, label_prefix: str = "附件"
) -> list[dict[str, str]]:
    """将逗号/JSON URL 列表转为 {url, label} 列表。"""
    out: list[dict[str, str]] = []
    for i, url in enumerate(parse_url_list(raw), start=1):
        out.append({"url": url, "label": f"{label_prefix}{i}"})
    return out


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 _plan_attachment_jobs(
    *,
    rows: Sequence[Mapping[str, Any]],
    attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
    id_field: str,
    attach_dir: str,
) -> Tuple[dict[str, list[tuple[str, str, str]]], list[tuple[str, str, str]]]:
    """预先确定附件 zip 路径与超链接（尚未下载）。

    返回:
      entity_file_meta: entity_id -> [(zip_path, label, formula_rel), ...]
      download_jobs: [(zip_path, label, url), ...] 按写入顺序
    """
    row_no_by_id = first_row_no_by_entity_id(rows, id_field=id_field)
    entity_file_meta: dict[str, list[tuple[str, str, str]]] = {}
    download_jobs: list[tuple[str, str, str]] = []
    planned = 0

    for entity_id, items in (attachments_by_id or {}).items():
        seq = row_no_by_id.get(entity_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
            planned += 1
            if planned > EXPORT_ZIP_MAX_FILES:
                raise ExportValidationError(
                    f"附件文件数超过 {EXPORT_ZIP_MAX_FILES}，请缩小筛选范围"
                )
            item_label = str(item.get("label") or "附件")
            filename = build_attachment_filename(seq, item_label, url)
            zip_path = f"{attach_dir}/{filename}"
            formula_rel = f"附件/{filename}"
            meta.append((zip_path, item_label, formula_rel))
            download_jobs.append((zip_path, item_label, url))
        if meta:
            entity_file_meta[entity_id] = meta
    return entity_file_meta, download_jobs


def _build_xlsx_bytes(
    *,
    db: Optional[Session],
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    entity_file_meta: Mapping[str, Sequence[tuple[str, str, str]]],
    id_field: str,
    sheet: str,
) -> bytes:
    max_att = 0
    for meta in entity_file_meta.values():
        max_att = max(max_att, len(meta))
    k = min(EXPORT_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)
        entity_id = str(data.get(id_field) or "").strip()
        for i in range(1, k + 1):
            data[_attach_field(i)] = ""
        if entity_id and entity_id not in seen_first:
            seen_first.add(entity_id)
            meta = list(entity_file_meta.get(entity_id) or [])
            for i, (_zip_path, item_label, formula_rel) in enumerate(meta[:k], start=1):
                data[_attach_field(i)] = build_hyperlink_formula(formula_rel, item_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 = sheet
    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)

    if sheet == "计划归档":
        from app.services.pl_plan_archive_export import append_archive_list_merge_sheet

        append_archive_list_merge_sheet(wb, all_columns, enriched_rows, resolver)

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


def build_attachment_zip_bytes(
    *,
    db: Optional[Session],
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
    stamp: str,
    package_label: str,
    id_field: str = "id",
    sheet_title: Optional[str] = None,
    xlsx_basename: Optional[str] = None,
    download_fn: Optional[Callable[[str], bytes]] = None,
) -> tuple[bytes, list[str]]:
    """打包 ZIP：xlsx + 附件/；返回 (zip_bytes, failure_messages)。"""
    failures: list[str] = []
    chunks = list(
        iter_attachment_zip_chunks(
            db=db,
            tenant_no=tenant_no,
            columns=columns,
            rows=rows,
            attachments_by_id=attachments_by_id,
            stamp=stamp,
            package_label=package_label,
            id_field=id_field,
            sheet_title=sheet_title,
            xlsx_basename=xlsx_basename,
            download_fn=download_fn,
            failures_out=failures,
        )
    )
    return b"".join(chunks), failures


def iter_attachment_zip_chunks(
    *,
    db: Optional[Session],
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
    stamp: str,
    package_label: str,
    id_field: str = "id",
    sheet_title: Optional[str] = None,
    xlsx_basename: Optional[str] = None,
    download_fn: Optional[Callable[[str], bytes]] = None,
    failures_out: Optional[list[str]] = None,
) -> Iterator[bytes]:
    """边下载边产出 ZIP 字节，便于代理层在长时间打包时保持连接。"""
    download = download_fn or _default_download
    label = sanitize_zip_segment(package_label) or "导出"
    sheet = (sheet_title or label)[:31] or "导出"
    xlsx_name = xlsx_basename or f"(总表){label}.xlsx"
    root = f"{label}({stamp})"
    attach_dir = f"{root}/附件"

    entity_file_meta, download_jobs = _plan_attachment_jobs(
        rows=rows,
        attachments_by_id=attachments_by_id,
        id_field=id_field,
        attach_dir=attach_dir,
    )
    xlsx_bytes = _build_xlsx_bytes(
        db=db,
        tenant_no=tenant_no,
        columns=columns,
        rows=rows,
        entity_file_meta=entity_file_meta,
        id_field=id_field,
        sheet=sheet,
    )

    failures: list[str] = []
    total_bytes = 0
    file_count = 0
    buf = io.BytesIO()
    sent = 0

    def flush() -> bytes:
        nonlocal sent
        data = buf.getvalue()
        if len(data) <= sent:
            return b""
        chunk = data[sent:]
        sent = len(data)
        return chunk

    with zipfile.ZipFile(buf, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr(f"{root}/{xlsx_name}", xlsx_bytes)
        chunk = flush()
        if chunk:
            yield chunk

        for zip_path, item_label, url in download_jobs:
            try:
                content = download(url)
            except Exception as exc:  # noqa: BLE001 — 单文件失败不阻断整包
                failures.append(f"{item_label}: {url} ({exc})")
                continue
            file_count += 1
            total_bytes += len(content)
            if file_count > EXPORT_ZIP_MAX_FILES:
                raise ExportValidationError(
                    f"附件文件数超过 {EXPORT_ZIP_MAX_FILES}，请缩小筛选范围"
                )
            if total_bytes > EXPORT_ZIP_MAX_BYTES:
                raise ExportValidationError("附件总体积过大，请缩小筛选范围后重试")
            zf.writestr(zip_path, content)
            chunk = flush()
            if chunk:
                yield chunk

        if failures:
            note = "以下附件下载失败，已跳过：\n" + "\n".join(failures)
            zf.writestr(f"{root}/导出说明.txt", note.encode("utf-8"))
            chunk = flush()
            if chunk:
                yield chunk

    # ZipFile 关闭后写入中央目录
    chunk = flush()
    if chunk:
        yield chunk

    if failures_out is not None:
        failures_out.extend(failures)


def build_attachment_zip_streaming_response(
    *,
    db: Session,
    tenant_no: str,
    columns: Sequence[ExportColumn],
    rows: Sequence[Mapping[str, Any]],
    attachments_by_id: Mapping[str, Sequence[Mapping[str, str]]],
    stamp: str,
    package_label: str,
    id_field: str = "id",
    sheet_title: Optional[str] = None,
    xlsx_basename: Optional[str] = None,
    download_fn: Optional[Callable[[str], bytes]] = None,
) -> StreamingResponse:
    """立即返回 StreamingResponse：先发响应头，再边下附件边推送 ZIP 字节。"""
    label = sanitize_zip_segment(package_label) or "导出"
    filename = f"{label}({stamp}).zip"
    encoded = quote(filename)
    disposition = f"attachment; filename={encoded}; filename*=UTF-8''{encoded}"

    def generate() -> Iterator[bytes]:
        yield from iter_attachment_zip_chunks(
            db=db,
            tenant_no=tenant_no,
            columns=columns,
            rows=rows,
            attachments_by_id=attachments_by_id,
            stamp=stamp,
            package_label=package_label,
            id_field=id_field,
            sheet_title=sheet_title,
            xlsx_basename=xlsx_basename,
            download_fn=download_fn,
        )

    return StreamingResponse(
        generate(),
        media_type="application/zip",
        headers={
            "Content-Disposition": disposition,
            # 避免中间代理整包缓冲导致长时间无字节
            "X-Accel-Buffering": "no",
        },
    )