book-mcp/services/document_service.py

"""教辅文档 (book_document) 业务层 — MongoDB 版本"""
import re
from pymongo.errors import DuplicateKeyError

from db import get_collection, COL_BOOK, COL_DOC, next_id, BookDocument
from db.models import _now


def _infer_doc_type(filename: str) -> str:
    """从文件名猜题型：课时作业 / 周末提升 / 阶段质量评估 / 综合评估"""
    n = filename
    if re.search(r'阶段质量评估', n):
        return '阶段质量评估'
    if re.search(r'综合.*评估|期末', n):
        return '综合评估'
    if re.search(r'周末提升', n):
        return '周末提升'
    if re.search(r'课时|梯级|高考真题集训', n):
        return '课时作业'
    return 'unknown'


def add_document(book_id: int, original_filename: str, original_path: str,
                 title: str | None = None, doc_type: str | None = None,
                 order_index: int | None = None) -> dict:
    """添加文档。同 (book_id, original_path) 幂等。"""
    col_doc = get_collection(COL_DOC)
    col_book = get_collection(COL_BOOK)

    # book 存在性
    if not col_book.find_one({'_id': book_id, 'is_deleted': 0}):
        return {'ok': False, 'error': f'book {book_id} not found'}

    # 1. 算 doc_id（稳定哈希）
    doc_id = BookDocument.make_doc_id(book_id, original_path)

    # 2. 幂等查
    existing = col_doc.find_one({'_id': doc_id})
    if existing:
        bd = BookDocument.from_mongo(existing)
        return {
            'ok': True, 'idempotent': True, 'doc_id': bd.doc_id,
            'order_index': bd.order_index, 'status': bd.status,
            'doc_type': bd.doc_type,
        }

    # 3. order_index（自动递增）
    if order_index is None:
        cur_max = col_doc.find({'book_id': book_id}) \
                       .sort('order_index', -1).limit(1)
        cur_max_list = list(cur_max)
        order_index = (cur_max_list[0]['order_index'] + 1) if cur_max_list else 1

    # 4. doc_type 推断
    dt = doc_type or _infer_doc_type(original_filename)

    # 5. 入库
    doc = BookDocument(
        book_id=book_id, doc_id=doc_id,
        original_filename=original_filename, original_path=original_path,
        title=title, doc_type=dt, order_index=order_index,
    )
    try:
        col_doc.insert_one(doc.to_mongo())
    except DuplicateKeyError:
        # 竞争：重查
        existing = col_doc.find_one({'_id': doc_id})
        bd = BookDocument.from_mongo(existing)
        return {
            'ok': True, 'idempotent': True, 'doc_id': bd.doc_id,
            'order_index': bd.order_index, 'status': bd.status,
            'doc_type': bd.doc_type,
        }
    return {
        'ok': True, 'idempotent': False, 'doc_id': doc_id,
        'order_index': order_index, 'status': 'pending', 'doc_type': dt,
    }


def get_document(doc_id: str, include_files: bool = True) -> dict | None:
    col = get_collection(COL_DOC)
    d = col.find_one({'_id': doc_id})
    if not d:
        return None
    bd = BookDocument.from_mongo(d).to_dict()
    if include_files:
        from .file_service import get_files
        bd['files'] = get_files(doc_id)
    return bd


def list_documents(book_id: int, status: str | None = None,
                   page: int = 1, page_size: int = 50) -> dict:
    col = get_collection(COL_DOC)
    q: dict = {'book_id': book_id}
    if status:
        q['status'] = status
    total = col.count_documents(q)
    cursor = col.find(q).sort('order_index', 1) \
                  .skip((page - 1) * page_size).limit(page_size)
    docs = [BookDocument.from_mongo(d).to_dict() for d in cursor]
    return {'documents': docs, 'total': total, 'page': page, 'page_size': page_size}


def _mark_file_ready(doc_id: str, file_type: str, url: str, size: int,
                     meta: dict | None) -> dict:
    """mark_html/json_ready 通用实现"""
    from .file_service import upsert_file
    from db import COL_FILE, COL_DOC, COL_BOOK
    from db.database import load_config
    from services.rustfs_service import build_key

    col_doc = get_collection(COL_DOC)
    doc = col_doc.find_one({'_id': doc_id})
    if not doc:
        return {'ok': False, 'error': f'doc {doc_id} not found'}

    bucket = load_config()['rustfs']['default_bucket']
    key = build_key(doc['book_id'], doc_id, file_type)
    f = upsert_file(
        doc_id=doc_id, file_type=file_type, object_key=key, object_bucket=bucket,
        size=size, url=url, meta=meta,
    )
    new_status = 'html_ready' if file_type == 'html' else 'json_ready'
    col_doc.update_one({'_id': doc_id}, {'$set': {
        'status': new_status, 'updated_at': _now(), 'error_msg': None,
    }})
    return {'ok': True, 'file_id': f['id'], 'object_key': key, 'url': url}


def mark_html_ready(doc_id: str, html_url: str, html_size: int,
                    stats: dict | None = None) -> dict:
    """标记文档 html 已就绪"""
    return _mark_file_ready(doc_id, 'html', html_url, html_size, stats)


def mark_json_ready(doc_id: str, json_url: str, json_size: int,
                    schema_version: str | None = None) -> dict:
    """标记文档 json 已就绪"""
    meta = {'schema_version': schema_version} if schema_version else None
    return _mark_file_ready(doc_id, 'json', json_url, json_size, meta)


def mark_failed(doc_id: str, error_msg: str, stage: str = 'html') -> dict:
    """标记文档失败"""
    col = get_collection(COL_DOC)
    r = col.update_one({'_id': doc_id}, {'$set': {
        'status': 'failed', 'error_msg': f'[{stage}] {error_msg}'[:1000],
        'updated_at': _now(),
    }})
    return {'ok': r.matched_count > 0}