AIRegulation-Demo-Test-Backend/app/api/routes/docs.py

"""文档管理 API"""

from fastapi import APIRouter, UploadFile, File, HTTPException, BackgroundTasks
import os
import uuid
from datetime import datetime
from typing import Optional

from app.schemas.doc import (
    DocumentUploadResponse,
    DocumentListResponse,
    DocumentInfo,
    ParseResponse,
    EmbedResponse,
    TaskStatusResponse,
)
from app.core.config import settings
from app.services.minio import minio_service
from app.services.database import db_service, init_db, DocStatus
from app.services.tasks import generate_task_id, task_manager, get_task_status
from app.workflows.document_workflow import (
    generate_doc_id,
    run_parse_workflow,
    run_embedding_workflow,
)
from app.utils.logger import logger


router = APIRouter(prefix="/docs", tags=["文档管理"])

# 启动时初始化数据库
init_db()


def get_content_type(filename: str) -> str:
    """根据文件扩展名获取 Content-Type"""
    ext = os.path.splitext(filename)[1].lower()
    content_types = {
        ".pdf": "application/pdf",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".doc": "application/msword",
        ".txt": "text/plain",
    }
    return content_types.get(ext, "application/octet-stream")


@router.post("/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    background_tasks: BackgroundTasks = None,
):
    """
    上传法规文档到 MinIO，并自动触发异步解析

    流程：
    1. 验证文件格式
    2. 生成文档ID
    3. 上传到 MinIO
    4. 创建数据库记录
    5. 触发异步解析任务（后续可替换为 RabbitMQ）
    """
    # 检查文件格式
    allowed_ext = [".pdf", ".docx", ".doc", ".txt"]
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in allowed_ext:
        raise HTTPException(400, f"Unsupported file format: {ext}")

    # 检查文件大小
    content = await file.read()
    max_size = 50 * 1024 * 1024  # 50MB
    if len(content) > max_size:
        raise HTTPException(400, f"File size exceeds limit: {max_size // 1024 // 1024}MB")

    # 生成文档ID
    doc_id = generate_doc_id()

    # 构建 MinIO 存储路径
    storage_filename = f"{doc_id}_{file.filename}"
    minio_path = f"documents/{storage_filename}"

    try:
        # 上传到 MinIO
        content_type = get_content_type(file.filename)
        minio_url = minio_service.upload_file(
            minio_path,
            content,
            content_type,
        )

        # 创建数据库记录
        doc = db_service.create_document(
            doc_id=doc_id,
            filename=storage_filename,
            original_name=file.filename,
            minio_path=minio_path,
            size=len(content),
        )

        logger.info(f"Document uploaded: {doc_id} - {file.filename}")

        # 触发异步解析任务
        parse_task_id = generate_task_id()
        db_service.create_parse_task(parse_task_id, doc_id)

        # 使用 asyncio 异步执行解析（后续替换为 RabbitMQ）
        background_tasks.add_task(
            run_parse_workflow_sync,
            parse_task_id,
            doc_id,
        )

        return DocumentUploadResponse(
            doc_id=doc_id,
            filename=file.filename,
            size=len(content),
            status="uploaded",
            parse_task_id=parse_task_id,
        )

    except Exception as e:
        logger.error(f"Upload failed: {e}")
        raise HTTPException(500, f"Upload failed: {str(e)}")


def run_parse_workflow_sync(task_id: str, doc_id: str):
    """同步包装器，用于 BackgroundTasks"""
    import asyncio
    asyncio.run(run_parse_workflow(task_id, doc_id))


@router.get("/list", response_model=DocumentListResponse)
async def list_documents():
    """获取已索引文档列表"""
    docs = db_service.list_documents()
    return DocumentListResponse(
        docs=[
            DocumentInfo(
                id=d.id,
                name=d.original_name,
                chunks=d.chunks,
                status=d.status,
                created_at=d.created_at,
            )
            for d in docs
        ]
    )


@router.get("/{doc_id}", response_model=DocumentInfo)
async def get_document(doc_id: str):
    """获取单个文档信息"""
    doc = db_service.get_document(doc_id)
    if not doc:
        raise HTTPException(404, "Document not found")

    return DocumentInfo(
        id=doc.id,
        name=doc.original_name,
        chunks=doc.chunks,
        status=doc.status,
        created_at=doc.created_at,
    )


@router.post("/parse/{doc_id}", response_model=ParseResponse)
async def parse_document(
    doc_id: str,
    background_tasks: BackgroundTasks = None,
):
    """
    手动触发文档解析（如果文档已上传但未解析）
    """
    doc = db_service.get_document(doc_id)
    if not doc:
        raise HTTPException(404, "Document not found")

    if doc.status not in [DocStatus.uploaded.value, DocStatus.failed.value]:
        raise HTTPException(400, f"Document status is {doc.status}, cannot parse")

    # 创建解析任务
    task_id = generate_task_id()
    db_service.create_parse_task(task_id, doc_id)

    # 异步执行
    background_tasks.add_task(
        run_parse_workflow_sync,
        task_id,
        doc_id,
    )

    return ParseResponse(
        doc_id=doc_id,
        task_id=task_id,
        status="parsing",
    )


@router.post("/embed/{doc_id}", response_model=EmbedResponse)
async def embed_document(
    doc_id: str,
    background_tasks: BackgroundTasks = None,
):
    """
    触发文档向量化（需要文档已解析）
    """
    doc = db_service.get_document(doc_id)
    if not doc:
        raise HTTPException(404, "Document not found")

    if doc.status != DocStatus.parsed.value:
        raise HTTPException(400, f"Document must be parsed first. Current status: {doc.status}")

    # 创建向量化任务
    task_id = generate_task_id()
    db_service.create_parse_task(task_id, doc_id)

    # 异步执行
    background_tasks.add_task(
        run_embedding_workflow_sync,
        task_id,
        doc_id,
    )

    return EmbedResponse(
        doc_id=doc_id,
        task_id=task_id,
        status="embedding",
    )


def run_embedding_workflow_sync(task_id: str, doc_id: str):
    """同步包装器，用于 BackgroundTasks"""
    import asyncio
    asyncio.run(run_embedding_workflow(task_id, doc_id))


@router.get("/task/{task_id}", response_model=TaskStatusResponse)
async def get_task_status_api(task_id: str):
    """获取任务状态"""
    status = get_task_status(task_id)
    if not status:
        # 检查数据库中的任务记录
        task = db_service.get_parse_task(task_id)
        if task:
            return TaskStatusResponse(
                task_id=task_id,
                status=task.status,
                progress=task.progress or 0,
                message=task.message,
            )
        raise HTTPException(404, "Task not found")

    return TaskStatusResponse(
        task_id=task_id,
        status=status.get("status", "unknown"),
        progress=status.get("progress", 0),
        message=status.get("message"),
        result=status.get("result"),
    )


@router.delete("/delete/{doc_id}")
async def delete_document(doc_id: str):
    """
    删除文档

    同时删除：
    - MinIO 中的文件
    - 数据库中的记录
    - 解析后的文本文件
    """
    doc = db_service.get_document(doc_id)
    if not doc:
        raise HTTPException(404, "Document not found")

    try:
        # 删除 MinIO 文件
        minio_service.delete_file(doc.minio_path)

        # 删除本地解析文件
        parsed_path = f"{settings.data_parsed_dir}/{doc_id}.txt"
        if os.path.exists(parsed_path):
            os.remove(parsed_path)

        # 删除本地临时文件
        temp_path = f"{settings.data_raw_dir}/{doc.filename}"
        if os.path.exists(temp_path):
            os.remove(temp_path)

        # 删除数据库记录
        db_service.delete_document(doc_id)

        logger.info(f"Document deleted: {doc_id}")

        return {"success": True, "doc_id": doc_id}

    except Exception as e:
        logger.error(f"Delete failed: {e}")
        raise HTTPException(500, f"Delete failed: {str(e)}")