doc-forge/backend/routers/export.py

import asyncio
import json
import os
import uuid
from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import PlainTextResponse, RedirectResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from config import settings
from database import get_db
from models.document import Document
from models.generation_log import GenerationLog
from models.paragraph import Paragraph
from models.template import Template
from models.template_block import TemplateBlock
from services.document_export import export_document_bytes
from services.minio_client import (
    download_object_bytes,
    get_presigned_url,
    split_bucket_path,
    upload_bytes,
)

router = APIRouter()


def _block_text_content(block: TemplateBlock) -> str:
    try:
        payload = json.loads(block.content_json or "{}")
    except Exception:
        payload = {}
    return payload.get("text") or block.default_value or ""


def _block_table_content(block: TemplateBlock) -> dict:
    try:
        payload = json.loads(block.content_json or "{}")
    except Exception:
        payload = {}
    return payload.get("table") or {}


def _build_block_export_content(block: TemplateBlock) -> dict:
    if block.block_type == "table":
        table_data = _block_table_content(block)
        matrix = table_data.get("data") or []
        headers = matrix[0] if matrix else []
        rows = matrix[1:] if len(matrix) > 1 else []
        return {"content": [{"type": "table", "headers": headers, "rows": rows}]}
    return {"content": [{"type": "text", "text": _block_text_content(block)}]}


def _resolve_block_write_modes(blocks: list[TemplateBlock]) -> list[str]:
    modes: list[str] = []
    anchor_counter: dict[str, int] = {}
    for block in blocks:
        if block.block_type == "heading":
            modes.append("replace_heading_only")
            continue
        anchor = (block.anchor_ref or block.title or "").strip()
        seen = anchor_counter.get(anchor, 0)
        modes.append("replace_section" if seen == 0 else "append_after_heading")
        anchor_counter[anchor] = seen + 1
    return modes


@router.get("/{document_id}/docx")
async def export_docx(document_id: int, db: AsyncSession = Depends(get_db)):
    document = await db.get(Document, document_id)
    if document is None:
        raise HTTPException(status_code=404, detail="生成记录不存在")

    template = await db.get(Template, document.template_id)
    if template is None:
        raise HTTPException(status_code=404, detail="模板不存在")

    template_bucket, template_object = split_bucket_path(template.file_path)
    template_bytes = await asyncio.to_thread(download_object_bytes, template_bucket, template_object)

    block_result = await db.execute(
        select(TemplateBlock)
        .where(TemplateBlock.template_id == template.id)
        .order_by(TemplateBlock.sort_index.asc(), TemplateBlock.id.asc())
    )
    blocks = block_result.scalars().all()

    log_result = await db.execute(
        select(GenerationLog).where(GenerationLog.document_id == document_id)
    )
    generation_logs = log_result.scalars().all()
    log_map = {item.paragraph_id: json.loads(item.content) if item.content else {"content": []} for item in generation_logs}

    logs = []
    if blocks:
        write_modes = _resolve_block_write_modes(blocks)
        for block, write_mode in zip(blocks, write_modes):
            generated_content = log_map.get(block.source_paragraph_id) if block.source_paragraph_id else None
            content = generated_content if (block.edit_mode == "ai" or block.block_type == "ai_slot") and generated_content else _build_block_export_content(block)
            logs.append(
                {
                    "anchor_title": block.anchor_ref or block.title,
                    "title": block.title,
                    "write_mode": write_mode,
                    "content": content,
                }
            )
    else:
        result = await db.execute(
            select(GenerationLog, Paragraph)
            .join(Paragraph, Paragraph.id == GenerationLog.paragraph_id)
            .where(GenerationLog.document_id == document_id)
            .order_by(Paragraph.sort_index.asc(), Paragraph.id.asc())
        )
        for log, paragraph in result.all():
            logs.append(
                {
                    "anchor_title": paragraph.anchor_title or paragraph.title,
                    "title": paragraph.title,
                    "write_mode": paragraph.write_mode,
                    "content": json.loads(log.content) if log.content else {"content": []},
                }
            )

    exported_bytes = await asyncio.to_thread(export_document_bytes, template_bytes, logs)
    object_name = f"{datetime.now().strftime('%Y%m%d')}/{uuid.uuid4().hex}.docx"
    await asyncio.to_thread(
        upload_bytes,
        settings.MINIO_BUCKET_OUTPUTS,
        object_name,
        exported_bytes,
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    )

    document.file_path = f"{settings.MINIO_BUCKET_OUTPUTS}/{object_name}"
    await db.commit()
    return RedirectResponse(
        url=get_presigned_url(settings.MINIO_BUCKET_OUTPUTS, object_name),
        status_code=307,
    )


@router.get("/{document_id}/pdf")
async def export_pdf(document_id: int):
    return PlainTextResponse(
        f"文档 {document_id} 的 PDF 导出功能正在开发中，当前版本请先使用预览页查看结果。",
        media_type="text/plain; charset=utf-8",
    )