app/services/ingestion.py

"""文档摄入编排器.

完整流水线:
1. 保存上传文件到 /data/uploads/{doc_id}/{filename}
2. 计算 SHA256, 查 SQLite 是否有同 sha 的 ready 文档
   - 有 -> 直接返回 {doc_id, status: "duplicate"}
3. 解析 (智能路由: primary -> fallback)
4. 分块 (层次化 + 可选语义 + 可选上下文预置)
5. 向量化 (BGE-M3 三路)
6. 入库:
   - ChromaDB: dense + (colbert)
   - 旁路: sparse (BM25-style)
   - SQLite: document 行 + chunk 行
7. 调度 persist.push 异步同步
8. 返回 {doc_id, chunk_count, status: "ready"}
"""
from __future__ import annotations

import asyncio
import hashlib
import logging
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from app.config import settings
from app.core.errors import IngestionFailedError
from app.core.paths import upload_dir
from app.models import db
from app.models.schemas import IngestResult
from app.services.chunking import (
    Chunk,
    chunk_document,
    chunk_to_embed_text,
    contextualize_chunks,
)
from app.services.embedding import get_embedder
from app.services.parsers import parse_with_fallback
from app.services.parsers.base_parser import ParsedDocument
from app.services.persist import push_to_hf, schedule_push
from app.services.vector_store import upsert_chunks

logger = logging.getLogger(__name__)


# 单并发摄入队列 (避免大 PDF 同时跑爆 RAM)
_ingest_lock = asyncio.Lock()


@dataclass
class _IngestContext:
    doc_id: str
    file_path: Path
    filename: str
    mime: str
    size: int
    sha256: str
    started: float


async def _save_upload(*, content: bytes, filename: str) -> tuple[Path, str, int, str]:
    """保存上传字节到 /data/uploads/{doc_id}/{filename}, 返回 (path, mime, size, sha256)."""
    doc_id = uuid.uuid4().hex
    dest_dir = upload_dir() / doc_id
    dest_dir.mkdir(parents=True, exist_ok=True)
    # 防路径穿越
    safe_name = Path(filename).name or "upload"
    dest = dest_dir / safe_name
    dest.write_bytes(content)

    # MIME 简单嗅探
    mime = _sniff_mime(safe_name, content[:2048])
    sha = hashlib.sha256(content).hexdigest()
    return dest, mime, len(content), sha


def _sniff_mime(filename: str, head: bytes) -> str:
    ext = Path(filename).suffix.lower()
    return {
        ".pdf": "application/pdf",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
        ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".tiff": "image/tiff",
        ".html": "text/html",
        ".md": "text/markdown",
        ".markdown": "text/markdown",
    }.get(ext, "application/octet-stream")


async def ingest_bytes(
    *,
    content: bytes,
    filename: str,
) -> IngestResult:
    """公开入口. 处理上传的原始字节."""
    async with _ingest_lock:
        return await _ingest_locked(content, filename)


async def _ingest_locked(content: bytes, filename: str) -> IngestResult:
    started = time.time()
    # 1. 保存
    file_path, mime, size, sha = await _save_upload(content=content, filename=filename)
    doc_id = file_path.parent.name  # 用 upload 子目录名当 doc_id (一致)

    logger.info("Ingest start: %s (%d bytes, sha=%s)", filename, size, sha[:12])

    # 2. SHA256 幂等检查
    existing = db.doc_find_by_sha256(sha)
    if existing is not None:
        logger.info("Duplicate by sha256: existing doc_id=%s", existing["id"])
        # 删除刚写入的副本, 复用旧 doc
        try:
            file_path.unlink(missing_ok=True)
        except Exception:  # noqa: BLE001
            pass
        return IngestResult(doc_id=existing["id"], chunk_count=existing.get("chunk_count", 0), status="duplicate")

    # 3. 写入 SQLite (状态 uploading -> parsing)
    db.doc_insert({
        "id": doc_id,
        "filename": Path(filename).name,
        "mime": mime,
        "size": size,
        "sha256": sha,
        "status": "parsing",
        "progress": 0,
        "progress_label": "开始解析...",
        "parser": settings.parser_primary,
        "chunk_size": settings.chunk_size,
        "chunk_overlap": settings.chunk_overlap,
        "semantic_chunking": settings.semantic_chunking,
        "contextual_retrieval": settings.contextual_retrieval,
        "created_at": started,
        "updated_at": started,
        "meta": {},
    })

    try:
        # 4. 解析
        db.doc_update_status(doc_id, "parsing", progress=10, progress_label="解析 PDF/Word 结构...")
        parsed = await parse_with_fallback(file_path)
        page_count = parsed.meta.get("page_count")
        db.doc_update_status(doc_id, "parsing", page_count=page_count, progress=30, progress_label="解析完成, 准备分块")

        # 5. 分块
        chunking = chunk_document(parsed, doc_id=doc_id)
        logger.info(
            "Chunked: %d parents, %d children for %s",
            len(chunking.parents), len(chunking.children), filename,
        )
        db.doc_update_status(
            doc_id, "parsing",
            progress=45,
            progress_label=f"分块 {len(chunking.children)} 个 chunk",
        )

        # 6. 上下文预置 (可选, 走 LLM, 慢但提升召回)
        if settings.contextual_retrieval and chunking.children:
            db.doc_update_status(
                doc_id, "embedding",
                progress=50,
                progress_label="上下文预标注 (走 LLM)...",
            )
            await contextualize_chunks(
                chunking.children, doc_title=Path(filename).name, max_concurrency=4
            )

        # 7. 向量化
        db.doc_update_status(doc_id, "embedding", progress=55, progress_label="向量化中...")
        embedder = get_embedder()
        embed_texts = [chunk_to_embed_text(c) for c in chunking.children]
        if not embed_texts:
            logger.warning("No chunks to embed for %s", filename)
            db.doc_update_status(doc_id, "ready", chunk_count=0, progress=100, progress_label="完成 (空文档)")
            return IngestResult(doc_id=doc_id, chunk_count=0, status="ready")

        # batch 编码
        BATCH = 32
        all_dense: list[np.ndarray] = []
        all_sparse: list[dict[int, float]] = []
        all_colbert: list[np.ndarray] = []
        n_total = len(embed_texts)
        for i in range(0, n_total, BATCH):
            batch = embed_texts[i : i + BATCH]
            out = await embedder.encode(batch)
            all_dense.append(out["dense"])
            all_sparse.extend(out.get("sparse", []))
            all_colbert.extend(out.get("colbert", []))
            done = min(i + BATCH, n_total)
            # 55% -> 95%
            pct = 55 + int((done / n_total) * 40)
            db.doc_update_status(
                doc_id, "embedding",
                progress=pct,
                progress_label=f"向量化 {done}/{n_total}",
            )
            logger.info(
                "Embedded %d/%d chunks", done, n_total
            )

        dense_arr = np.vstack(all_dense) if all_dense else np.zeros((0, 1024), dtype=np.float32)

        # 8. 入库 ChromaDB + sparse 旁路
        ids = [c.id for c in chunking.children]
        documents = [c.text for c in chunking.children]
        metadatas = [
            {
                "doc_id": c.doc_id,
                "parent_id": c.parent_id or "",
                "chunk_index": c.chunk_index,
                "page_no": c.page_no or 0,
                "heading": c.heading or "",
                "context_prefix": c.context_prefix or "",
            }
            for c in chunking.children
        ]
        upsert_chunks(
            ids=ids,
            embeddings=dense_arr,
            documents=documents,
            metadatas=metadatas,
            sparse_weights=all_sparse or None,
            colbert_vecs=all_colbert or None,
        )

        # 9. 写入 SQLite chunks
        db.chunk_insert_bulk([
            {
                "id": c.id,
                "doc_id": c.doc_id,
                "parent_id": c.parent_id,
                "chunk_index": c.chunk_index,
                "text": c.text,
                "token_count": c.token_count,
                "page_no": c.page_no,
                "heading": c.heading,
                "context_prefix": c.context_prefix,
                "created_at": c.created_at,
            }
            for c in chunking.children
        ])

        # 10. 文档就绪
        elapsed_ms = int((time.time() - started) * 1000)
        db.doc_update_status(
            doc_id, "ready",
            chunk_count=len(chunking.children),
            page_count=page_count,
            progress=100,
            progress_label="就绪",
        )
        logger.info(
            "Ingest done: %s doc_id=%s chunks=%d %dms",
            filename, doc_id, len(chunking.children), elapsed_ms,
        )

        # 11. ✅ A 改良版: 重要写操作 (upload) 同步阻塞 push, 不再吃 debounce 黑洞
        # 用户已经等了 1-3 分钟摄入, 再加 5-15s push 可接受
        # push 内部已带 verify, 失败会写 _state['last_error'] 让 /readyz 暴露
        push_ok = await push_to_hf()
        if not push_ok:
            # push 失败但不阻塞 ingest 结果 (用户至少能看到 doc ready)
            logger.warning(
                "Ingest done but persist push failed for %s. "
                "See /readyz persist.last_error. "
                "Data is still on local /data but may not survive Space restart.",
                doc_id,
            )

        return IngestResult(
            doc_id=doc_id,
            chunk_count=len(chunking.children),
            status="ready",
        )

    except Exception as e:  # noqa: BLE001
        logger.exception("Ingest failed for %s", filename)
        try:
            db.doc_update_status(doc_id, "failed", error=str(e)[:500])
        except Exception:  # noqa: BLE001
            pass
        raise IngestionFailedError(
            f"Ingest failed: {e}", code="ingest_failed", detail={"doc_id": doc_id}
        ) from e


async def delete_document(doc_id: str) -> bool:
    """完整删除: ChromaDB + 旁路 sparse + SQLite + uploads/."""
    # 先查, 拿文件路径
    doc = db.doc_get(doc_id)
    if doc is None:
        return False

    # 从 ChromaDB / sparse 旁路删
    try:
        from app.services.vector_store import delete_by_doc
        delete_by_doc(doc_id)
    except Exception as e:  # noqa: BLE001
        logger.warning("ChromaDB delete failed for %s: %s", doc_id, e)

    # SQLite chunks + document
    db.chunk_delete_by_doc(doc_id)
    db.doc_delete(doc_id)

    # 删上传原文件
    if doc.get("sha256"):
        # doc_id 即子目录名
        up_dir = upload_dir() / doc_id
        if up_dir.exists():
            import shutil
            shutil.rmtree(up_dir, ignore_errors=True)

    await schedule_push()
    # ⚠️ 注意: delete 仍走 schedule_push (debounce), 不阻塞用户响应.
    # 因为 delete 操作不产生新数据, 延迟 5s 推完全可接受.
    return True


def list_documents(limit: int = 100) -> list[dict[str, Any]]:
    return db.doc_list(limit=limit)


def get_document(doc_id: str) -> dict[str, Any] | None:
    return db.doc_get(doc_id)