plugins/doc/api.py

import logging
import re
import tempfile
import zipfile
from pathlib import Path
from typing import List

from fastapi import APIRouter, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse

logger = logging.getLogger(__name__)

router = APIRouter()

plugin = None


def set_plugin_instance(plugin_instance):
    global plugin
    plugin = plugin_instance


@router.get("/status")
async def get_status():
    if plugin is None:
        return {
            "name": "doc",
            "enabled": False,
            "message": "插件未加载",
        }
    return plugin.get_status()


@router.post("/merge")
async def merge_documents(
    files: List[UploadFile] = File(...),
    order: str = Form(None),
    output_format: str = Form("md"),
):
    if plugin is None or not plugin.enabled:
        raise HTTPException(status_code=400, detail="插件未启用")

    if output_format not in {"md", "txt"}:
        raise HTTPException(status_code=400, detail="格式参数错误，仅支持 md 或 txt")

    if not files:
        raise HTTPException(status_code=400, detail="至少需要上传一个文件")

    allowed_extensions = {".md", ".txt"}
    contents = []

    for file in files:
        ext = Path(file.filename).suffix.lower()
        if ext not in allowed_extensions:
            raise HTTPException(
                status_code=400, detail=f"不支持的文件格式: {file.filename}"
            )

        content = await file.read()
        try:
            decoded_content = content.decode("utf-8")
            contents.append(decoded_content)
        except UnicodeDecodeError:
            raise HTTPException(
                status_code=400, detail=f"文件编码不支持: {file.filename}"
            )

    if order:
        try:
            order_list = [int(x) for x in order.split(",")]
            reordered_contents = [
                contents[i] for i in order_list if 0 <= i < len(contents)
            ]
            if not reordered_contents:
                raise HTTPException(status_code=400, detail="顺序参数无效")
            contents = reordered_contents
        except (ValueError, IndexError):
            raise HTTPException(
                status_code=400, detail="顺序参数格式错误，应为逗号分隔的数字"
            )

    merged_content = "\n\n".join([item.strip() for item in contents])

    file_suffix = f".{output_format}"
    temp_file = tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=file_suffix, encoding="utf-8"
    )
    temp_file.write(merged_content)
    temp_file.close()

    media_type = "text/markdown" if output_format == "md" else "text/plain"
    filename = f"merged_document.{output_format}"

    return FileResponse(temp_file.name, media_type=media_type, filename=filename)


@router.post("/split")
async def split_document(file: UploadFile = File(...), output_format: str = Form("md")):
    if plugin is None or not plugin.enabled:
        raise HTTPException(status_code=400, detail="插件未启用")

    if output_format not in {"md", "txt"}:
        raise HTTPException(status_code=400, detail="格式参数错误，仅支持 md 或 txt")

    ext = Path(file.filename).suffix.lower()
    if ext not in {".md", ".txt"}:
        raise HTTPException(
            status_code=400, detail="不支持的文件格式，仅支持 .md 和 .txt"
        )

    content_bytes = await file.read()
    try:
        content = content_bytes.decode("utf-8")
    except UnicodeDecodeError:
        raise HTTPException(status_code=400, detail="文件编码不支持，请使用 UTF-8 编码")

    pattern = re.compile(r"^第\d+章", re.MULTILINE)
    matches = list(pattern.finditer(content))

    if not matches:
        raise HTTPException(status_code=400, detail="未找到 '第X章' 格式的章节标记")

    temp_dir = tempfile.mkdtemp()
    zip_path = Path(temp_dir) / "split_documents.zip"

    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for i, match in enumerate(matches):
            start_pos = match.start()
            end_pos = matches[i + 1].start() if i + 1 < len(matches) else len(content)

            chapter_content = content[start_pos:end_pos].strip()
            first_line = chapter_content.split("\n", 1)[0].strip()

            sanitized = re.sub(r'[<>:"/\\|?*]', "_", first_line)
            filename = f"{sanitized}.{output_format}"

            temp_file = Path(temp_dir) / filename
            with open(temp_file, "w", encoding="utf-8") as f:
                f.write(chapter_content)

            zipf.write(temp_file, filename)

    return FileResponse(
        zip_path, media_type="application/zip", filename="split_documents.zip"
    )