vent_ai_platform/app/api/v1/endpoints/documents.py

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from urllib.parse import quote

from fastapi import (
    APIRouter,
    BackgroundTasks,
    Depends,
    File,
    Form,
    Query,
    UploadFile,
    status,
)
from fastapi.responses import FileResponse, StreamingResponse
from sqlalchemy.orm import Session

from app.core.config import PROJECT_ROOT, settings
from app.core.errors import APIError
from app.db.repositories import document as document_repo
from app.db.repositories import mine as mine_repo
from app.db.session import get_db
from app.schemas.common import MessageResponse, Page
from app.schemas.document import (
    BusinessDocumentCreate,
    BusinessDocumentDetail,
    BusinessDocumentRead,
    BusinessDocumentStatistics,
    DocumentPreviewRead,
    BusinessDocumentUpdate,
    DocumentCategorySummary,
    DocumentExtractionCreate,
    DocumentExtractionRead,
    DocumentExtractionUpdate,
    DocumentExtractionTemplateRead,
    DocumentProcessingRunRead,
    DocumentProcessingStart,
    DocumentProcessingTaskRead,
    DocumentSearchResponse,
)
from app.services.document_service import (
    add_extraction,
    create_document,
    delete_document,
    get_document_detail,
    process_extraction_edit_judgement,
    update_document,
    update_extraction,
)
from app.services.document_processing import (
    process_document_run,
    start_document_processing,
)
from app.services.document_business_templates import resolve_document_template
from app.services.document_export import build_document_list_workbook
from app.services.storage_service import document_storage_statistics


router = APIRouter(prefix="/documents", tags=["业务资料"])
MAX_UPLOAD_BYTES = 50 * 1024 * 1024
MIME_TYPES = {
    ".pdf": "application/pdf",
    ".doc": "application/msword",
    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ".xls": "application/vnd.ms-excel",
    ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ".csv": "text/csv",
}


def _resolve_document_file_path(document) -> Path:
    source = Path(document.file_path)
    if not source.is_absolute():
        source = PROJECT_ROOT / source
    resolved = source.resolve()
    upload_root = settings.upload_path.resolve()
    if not resolved.is_relative_to(upload_root):
        raise APIError(403, "document_file_forbidden", "资料文件路径不允许访问")
    return resolved


@router.get("", response_model=Page[BusinessDocumentRead], summary="资料列表")
def list_documents(
    mine_id: str = Query(...),
    category: str | None = Query(None),
    keyword: str | None = Query(None),
    parse_status: str | None = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=200),
    db: Session = Depends(get_db),
) -> Page[BusinessDocumentRead]:
    items, total = document_repo.list_documents(
        db,
        mine_id=mine_id,
        category=category,
        keyword=keyword,
        parse_status=parse_status,
        page=page,
        page_size=page_size,
    )
    return Page(items=items, total=total, page=page, page_size=page_size)


@router.get(
    "/categories",
    response_model=list[DocumentCategorySummary],
    summary="资料分类统计",
)
def category_summary(
    mine_id: str = Query(...),
    db: Session = Depends(get_db),
) -> list[DocumentCategorySummary]:
    return [
        DocumentCategorySummary(category=category, count=count)
        for category, count in document_repo.category_summary(db, mine_id)
    ]


@router.get(
    "/statistics",
    response_model=BusinessDocumentStatistics,
    summary="业务资料抽取统计",
)
def statistics(
    mine_id: str = Query(...),
    db: Session = Depends(get_db),
) -> BusinessDocumentStatistics:
    statistics_data = document_repo.document_statistics(db, mine_id)
    statistics_data.update(document_storage_statistics())
    return BusinessDocumentStatistics(**statistics_data)


@router.get(
    "/search",
    response_model=DocumentSearchResponse,
    summary="全资料库全文检索",
)
def search_documents(
    q: str = Query(..., min_length=1, max_length=200),
    mine_id: str = Query(...),
    category: str | None = Query(None),
    page: int = Query(1, ge=1),
    page_size: int = Query(10, ge=1, le=50),
    db: Session = Depends(get_db),
) -> DocumentSearchResponse:
    hits, total = document_repo.search_documents(
        db,
        mine_id=mine_id,
        query=q,
        category=category,
        page=page,
        page_size=page_size,
    )
    return DocumentSearchResponse(
        query=q,
        total=total,
        page=page,
        page_size=page_size,
        items=hits,
    )


@router.get(
    "/processing-tasks",
    response_model=list[DocumentProcessingTaskRead],
    summary="最近业务资料处理任务",
)
def processing_tasks(
    mine_id: str = Query(...),
    limit: int = Query(8, ge=1, le=100),
    db: Session = Depends(get_db),
) -> list[DocumentProcessingTaskRead]:
    return [
        DocumentProcessingTaskRead(**item)
        for item in document_repo.list_processing_tasks(
            db,
            mine_id=mine_id,
            limit=limit,
        )
    ]


@router.get(
    "/templates",
    response_model=list[DocumentExtractionTemplateRead],
    summary="资料抽取模板",
)
def extraction_templates(
    document_type: str | None = Query(None),
    db: Session = Depends(get_db),
) -> list[DocumentExtractionTemplateRead]:
    return document_repo.list_templates(db, document_type=document_type)


@router.get(
    "/export",
    summary="导出业务资料清单",
)
def export_documents(
    mine_id: str = Query(...),
    category: str | None = Query(None),
    keyword: str | None = Query(None),
    parse_status: str | None = Query(None),
    db: Session = Depends(get_db),
) -> StreamingResponse:
    documents = document_repo.list_documents_for_export(
        db,
        mine_id=mine_id,
        category=category,
        keyword=keyword,
        parse_status=parse_status,
    )
    counts = document_repo.extraction_counts(
        db,
        [document.id for document in documents],
    )
    mine = mine_repo.get_mine(db, mine_id)
    mine_name = mine.name if mine else "矿井"
    category_label = category if category and category != "all" else "全部资料"
    document_type_labels: dict[str, str] = {}
    for template in document_repo.list_templates(db):
        document_type_labels[template.code] = template.name
        if template.document_type:
            document_type_labels.setdefault(template.document_type, template.name)
    workbook = build_document_list_workbook(
        documents,
        counts,
        mine_name=mine_name,
        category_label=category_label,
        keyword=keyword or "",
        document_type_labels=document_type_labels,
    )
    timestamp = datetime.now().strftime("%Y%m%d_%H%M")
    filename = f"{mine_name}业务资料清单_{timestamp}.xlsx"
    headers = {
        "Content-Disposition": (
            f"attachment; filename=document_list_{timestamp}.xlsx; "
            f"filename*=UTF-8''{quote(filename)}"
        )
    }
    return StreamingResponse(
        workbook,
        media_type=(
            "application/vnd.openxmlformats-officedocument."
            "spreadsheetml.sheet"
        ),
        headers=headers,
    )


@router.post(
    "",
    response_model=BusinessDocumentRead,
    status_code=status.HTTP_201_CREATED,
    summary="创建资料元数据",
)
def create_document_record(
    payload: BusinessDocumentCreate,
    db: Session = Depends(get_db),
):
    return create_document(db, payload)


@router.post(
    "/upload",
    response_model=BusinessDocumentRead,
    status_code=status.HTTP_201_CREATED,
    summary="上传业务资料",
)
async def upload_document(
    background_tasks: BackgroundTasks,
    mine_id: str = Form(...),
    title: str | None = Form(None),
    category: str = Form(...),
    owner_name: str = Form(""),
    department: str = Form(""),
    code: str | None = Form(None),
    parse_mode: str = Form("lightweight"),
    template_id: str | None = Form(None),
    file: UploadFile = File(...),
    db: Session = Depends(get_db),
):
    settings.upload_path.mkdir(parents=True, exist_ok=True)
    safe_name = Path(file.filename or "document.bin").name
    resolved_title = (title or Path(safe_name).stem).strip() or "未命名资料"
    template = resolve_document_template(
        db,
        title=resolved_title,
        template_id=template_id,
    )
    target = settings.upload_path / f"{mine_id}_{safe_name}"
    content = await file.read(MAX_UPLOAD_BYTES + 1)
    await file.close()
    if len(content) > MAX_UPLOAD_BYTES:
        raise APIError(413, "file_too_large", "单个资料文件不能超过 50 MiB")
    target.write_bytes(content)
    payload = BusinessDocumentCreate(
        mine_id=mine_id,
        code=code,
        title=resolved_title,
        category=category,
        file_name=safe_name,
        file_path=str(target),
        file_ext=Path(safe_name).suffix.lower(),
        file_size=len(content),
        owner_name=owner_name,
        department=department,
    )
    document = create_document(db, payload)
    run = start_document_processing(
        db,
        document.id,
        parse_mode=parse_mode,
        template_id=template.id,
    )
    background_tasks.add_task(process_document_run, run.id)
    return document


@router.get(
    "/{document_id}",
    response_model=BusinessDocumentDetail,
    summary="资料详情",
)
def document_detail(document_id: str, db: Session = Depends(get_db)):
    document, extractions = get_document_detail(db, document_id)
    document_data = BusinessDocumentRead.model_validate(document).model_dump()
    document_data["metadata_json"] = document_data.pop("metadata")
    return BusinessDocumentDetail(
        **document_data,
        extractions=extractions,
    )


@router.get(
    "/{document_id}/preview",
    response_model=DocumentPreviewRead,
    summary="资料预览",
)
def document_preview(
    document_id: str,
    db: Session = Depends(get_db),
) -> DocumentPreviewRead:
    document = document_repo.get_document(db, document_id)
    if document is None:
        raise APIError(404, "document_not_found", "业务资料不存在")
    source = _resolve_document_file_path(document)
    source_available = source.is_file()
    latest_run = document_repo.get_latest_processing_run(db, document_id)
    parsed_text = str(latest_run.parsed_markdown or "") if latest_run else ""
    extension = (document.file_ext or source.suffix).lower()
    if not parsed_text and source_available and extension == ".csv":
        try:
            parsed_text = source.read_text(encoding="utf-8-sig", errors="replace")
        except OSError:
            parsed_text = ""
    text_available = bool(parsed_text.strip())
    if source_available and extension == ".pdf":
        preview_type = "pdf"
    elif text_available:
        preview_type = "text"
    elif source_available:
        preview_type = "download_only"
    else:
        preview_type = "unavailable"
    base_url = f"{settings.api_v1_prefix}/documents/{document_id}"
    return DocumentPreviewRead(
        id=document.id,
        title=document.title,
        file_name=document.file_name,
        file_ext=extension,
        file_size=int(document.file_size or 0),
        preview_type=preview_type,
        source_available=source_available,
        text_available=text_available,
        can_download=source_available,
        file_url=f"{base_url}/file",
        download_url=f"{base_url}/file?download=1",
        text=parsed_text[:200000],
    )


@router.get(
    "/{document_id}/file",
    summary="原文件预览或下载",
)
def document_file(
    document_id: str,
    download: bool = Query(False),
    db: Session = Depends(get_db),
) -> FileResponse:
    document = document_repo.get_document(db, document_id)
    if document is None:
        raise APIError(404, "document_not_found", "业务资料不存在")
    source = _resolve_document_file_path(document)
    if not source.is_file():
        raise APIError(404, "document_file_not_found", "原始文件不存在")
    filename = document.file_name or source.name
    suffix = source.suffix.lower()
    disposition = "attachment" if download else "inline"
    headers = {
        "Content-Disposition": (
            f"{disposition}; filename=document{suffix}; "
            f"filename*=UTF-8''{quote(filename)}"
        )
    }
    return FileResponse(
        source,
        media_type=MIME_TYPES.get(suffix, "application/octet-stream"),
        headers=headers,
    )


@router.patch(
    "/{document_id}",
    response_model=BusinessDocumentRead,
    summary="更新资料",
)
def patch_document(
    document_id: str,
    payload: BusinessDocumentUpdate,
    db: Session = Depends(get_db),
):
    return update_document(db, document_id, payload)


@router.delete(
    "/{document_id}",
    response_model=MessageResponse,
    summary="删除资料",
)
def remove_document(document_id: str, db: Session = Depends(get_db)) -> MessageResponse:
    delete_document(db, document_id)
    return MessageResponse(message="资料已删除")


@router.post(
    "/{document_id}/process",
    response_model=DocumentProcessingRunRead,
    status_code=status.HTTP_202_ACCEPTED,
    summary="重新解析并抽取资料",
)
def process_document(
    document_id: str,
    payload: DocumentProcessingStart,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
):
    run = start_document_processing(
        db,
        document_id,
        force=payload.force,
        parse_mode=payload.parse_mode,
        template_id=payload.template_id,
    )
    background_tasks.add_task(process_document_run, run.id)
    return run


@router.get(
    "/{document_id}/processing",
    response_model=DocumentProcessingRunRead,
    summary="资料处理状态",
)
def processing_status(
    document_id: str,
    db: Session = Depends(get_db),
):
    run = document_repo.get_latest_processing_run(db, document_id)
    if run is None:
        raise APIError(404, "processing_run_not_found", "资料尚未开始解析")
    return run


@router.post(
    "/{document_id}/extractions",
    response_model=DocumentExtractionRead,
    status_code=status.HTTP_201_CREATED,
    summary="新增结构化抽取字段",
)
def create_extraction(
    document_id: str,
    payload: DocumentExtractionCreate,
    db: Session = Depends(get_db),
):
    if payload.document_id != document_id:
        payload.document_id = document_id
    return add_extraction(db, payload)


@router.patch(
    "/{document_id}/extractions/{extraction_id}",
    response_model=DocumentExtractionRead,
    summary="更新结构化抽取字段",
)
async def patch_extraction(
    document_id: str,
    extraction_id: str,
    payload: DocumentExtractionUpdate,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
):
    extraction, judgement_task = update_extraction(
        db,
        document_id,
        extraction_id,
        payload,
    )
    if judgement_task is not None:
        background_tasks.add_task(
            process_extraction_edit_judgement,
            judgement_task,
        )
    return extraction