vent_ai_platform/app/services/document_export.py

from __future__ import annotations

from datetime import datetime
from io import BytesIO
from typing import Any

from openpyxl import Workbook
from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
from openpyxl.utils import get_column_letter

from app.db.models.document import BusinessDocument


PARSE_STATUS_LABELS = {
    "pending": "等待解析",
    "running": "处理中",
    "parsing": "解析中",
    "extracting": "抽取中",
    "completed": "抽取完成",
    "failed": "处理失败",
}


def _safe_excel_value(value: Any) -> Any:
    if isinstance(value, str) and value.startswith(("=", "+", "-", "@")):
        return f"'{value}"
    return value


def build_document_list_workbook(
    documents: list[BusinessDocument],
    extraction_counts: dict[str, int],
    *,
    mine_name: str,
    category_label: str = "全部资料",
    keyword: str = "",
    document_type_labels: dict[str, str] | None = None,
) -> BytesIO:
    type_labels = document_type_labels or {}
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = "资料清单"
    headers = [
        "序号",
        "资料名称",
        "文档编号",
        "资料分组",
        "资料子类",
        "资料类型",
        "文件名",
        "扩展名",
        "文件大小(MB)",
        "页数",
        "行数",
        "资料责任人",
        "部门",
        "版本",
        "解析状态",
        "解析进度",
        "抽取字段数",
        "入库日期",
        "生效日期",
        "到期日期",
        "摘要",
    ]
    last_column = get_column_letter(len(headers))
    sheet.merge_cells(f"A1:{last_column}1")
    sheet["A1"] = f"{mine_name}业务资料清单"
    sheet["A1"].font = Font(size=16, bold=True, color="16222F")
    sheet["A1"].alignment = Alignment(horizontal="center", vertical="center")
    sheet.row_dimensions[1].height = 30

    sheet.merge_cells(f"A2:{last_column}2")
    filter_text = f"导出时间：{datetime.now().strftime('%Y-%m-%d %H:%M')}    范围：{category_label}"
    if keyword:
        filter_text += f"    搜索：{keyword}"
    sheet["A2"] = filter_text
    sheet["A2"].font = Font(size=10, color="6F7F8E")
    sheet["A2"].alignment = Alignment(horizontal="left", vertical="center")

    for column_index, header in enumerate(headers, start=1):
        cell = sheet.cell(row=3, column=column_index, value=header)
        cell.fill = PatternFill("solid", fgColor="0F9F8F")
        cell.font = Font(color="FFFFFF", bold=True)
        cell.alignment = Alignment(horizontal="center", vertical="center")
        cell.border = Border(
            bottom=Side(style="thin", color="D5E2E6"),
        )
    sheet.row_dimensions[3].height = 24

    date_format = "yyyy-mm-dd"
    thin_border = Border(
        bottom=Side(style="hair", color="E2E9EC"),
    )
    for row_index, document in enumerate(documents, start=4):
        values = [
            row_index - 3,
            document.title,
            document.code or "",
            document.category,
            document.subcategory or "未填写",
            type_labels.get(document.document_type, document.document_type or "未分类"),
            document.file_name,
            document.file_ext,
            round(float(document.file_size or 0) / (1024 * 1024), 2),
            document.page_count,
            document.row_count,
            document.owner_name,
            document.department,
            document.version,
            PARSE_STATUS_LABELS.get(document.parse_status, document.parse_status),
            f"{int(document.parse_progress or 0)}%",
            int(extraction_counts.get(document.id, 0)),
            document.created_at,
            document.effective_date,
            document.expire_date,
            document.summary,
        ]
        for column_index, value in enumerate(values, start=1):
            cell = sheet.cell(
                row=row_index,
                column=column_index,
                value=_safe_excel_value(value),
            )
            cell.alignment = Alignment(
                vertical="top",
                wrap_text=column_index in {2, 6, 7, 21},
            )
            cell.border = thin_border
            if column_index in {18, 19, 20} and value is not None:
                cell.number_format = date_format
            if column_index in {9, 10, 11, 16, 17}:
                cell.alignment = Alignment(horizontal="center", vertical="top")

    sheet.freeze_panes = "A4"
    sheet.auto_filter.ref = f"A3:{last_column}{max(sheet.max_row, 3)}"
    widths = {
        1: 7,
        2: 38,
        3: 20,
        4: 14,
        5: 16,
        6: 24,
        7: 34,
        8: 10,
        9: 14,
        10: 9,
        11: 9,
        12: 14,
        13: 14,
        14: 10,
        15: 13,
        16: 11,
        17: 13,
        18: 13,
        19: 13,
        20: 13,
        21: 48,
    }
    for column_index, width in widths.items():
        sheet.column_dimensions[get_column_letter(column_index)].width = width

    workbook.properties.title = f"{mine_name}业务资料清单"
    workbook.properties.creator = "矿安通防智能管控平台"
    output = BytesIO()
    workbook.save(output)
    workbook.close()
    output.seek(0)
    return output