doc-forge-reds/backend/app/api/templates.py

import uuid
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Query
from fastapi.responses import Response
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from app.core.database import get_db
from app.core.security_middleware import validate_file_extension
from app.models.template import Template
from app.schemas.template import TemplateResponse, TemplateListItem, HTMLContentResponse, HTMLUpdateRequest
from app.services.file_storage import save_upload, get_file_content, delete_file, TEMPLATES_DIR
from app.services.document_processor import docx_to_html, html_to_docx_bytes, docx_to_pdf_bytes

ALLOWED_TAGS = [
    "p", "div", "span", "br", "hr",
    "h1", "h2", "h3", "h4", "h5", "h6",
    "ul", "ol", "li",
    "a", "img", "table", "thead", "tbody", "tr", "td", "th",
    "b", "i", "u", "strong", "em", "del", "sub", "sup",
    "pre", "code", "blockquote",
]
ALLOWED_ATTRS = {
    "a": ["href", "title", "target"],
    "img": ["src", "alt", "width", "height"],
    "td": ["colspan", "rowspan"],
    "th": ["colspan", "rowspan"],
    "p": ["style"],
    "span": ["style"],
    "div": ["style"],
    "table": ["style"],
}

import bleach

router = APIRouter(prefix="/templates", tags=["模板管理"])


@router.post("", response_model=TemplateResponse)
async def upload_template(
    file: UploadFile = File(...),
    name: str | None = Form(None),
    db: AsyncSession = Depends(get_db),
):
    if not file.filename or not file.filename.endswith(".docx"):
        raise HTTPException(status_code=400, detail="仅支持 .docx 文件")

    validate_file_extension(file.filename)
    file_path = await save_upload(file, TEMPLATES_DIR)
    file_content = await get_file_content(file_path)
    html_content = await docx_to_html(file_content)

    template = Template(
        name=name or file.filename.replace(".docx", ""),
        file_path=file_path,
        html_content=html_content,
    )
    db.add(template)
    await db.flush()
    await db.refresh(template)
    return template


@router.get("", response_model=list[TemplateListItem])
async def list_templates(
    skip: int = Query(0, ge=0),
    limit: int = Query(20, ge=1, le=100),
    db: AsyncSession = Depends(get_db),
):
    query = select(Template).offset(skip).limit(limit).order_by(Template.created_at.desc())
    result = await db.execute(query)
    return result.scalars().all()


@router.get("/{template_id}", response_model=TemplateResponse)
async def get_template(template_id: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Template).where(Template.id == template_id))
    template = result.scalar_one_or_none()
    if not template:
        raise HTTPException(status_code=404, detail="模板不存在")
    return template


@router.get("/{template_id}/html", response_model=HTMLContentResponse)
async def get_template_html(template_id: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Template).where(Template.id == template_id))
    template = result.scalar_one_or_none()
    if not template:
        raise HTTPException(status_code=404, detail="模板不存在")
    return HTMLContentResponse(html_content=template.html_content or "")


@router.put("/{template_id}/html")
async def update_template_html(
    template_id: str,
    data: HTMLUpdateRequest,
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(select(Template).where(Template.id == template_id))
    template = result.scalar_one_or_none()
    if not template:
        raise HTTPException(status_code=404, detail="模板不存在")

    template.html_content = bleach.clean(
        data.html_content, tags=ALLOWED_TAGS, attributes=ALLOWED_ATTRS, strip=True
    )

    docx_bytes = html_to_docx_bytes(template.html_content)
    with open(template.file_path, "wb") as f:
        f.write(docx_bytes)

    await db.flush()
    return {"detail": "保存成功"}


@router.get("/{template_id}/download")
async def download_template(
    template_id: str,
    format: str = Query("docx", pattern="^(docx|pdf)$"),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(select(Template).where(Template.id == template_id))
    template = result.scalar_one_or_none()
    if not template:
        raise HTTPException(status_code=404, detail="模板不存在")

    file_content = await get_file_content(template.file_path)

    if format == "pdf":
        file_content = docx_to_pdf_bytes(file_content)
        media_type = "application/pdf"
        filename = f"{template.name}.pdf"
    else:
        media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        filename = f"{template.name}.docx"

    return Response(
        content=file_content,
        media_type=media_type,
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@router.delete("/{template_id}")
async def delete_template(template_id: str, db: AsyncSession = Depends(get_db)):
    result = await db.execute(select(Template).where(Template.id == template_id))
    template = result.scalar_one_or_none()
    if not template:
        raise HTTPException(status_code=404, detail="模板不存在")

    delete_file(template.file_path)
    await db.delete(template)
    return {"detail": "删除成功"}