"""标准 RAG 管线 (Phase 2.C 标准化)
====================================
设计基准（标准对照表见 STANDARDS_AUDIT_2026-07-24.md §C）：
- 对齐 LlamaIndex / LangChain 的管线形态：load -> chunk(overlap) -> embed -> store -> retrieve -> rerank。
- 复用现有 pgvector 后端（document_embeddings 表 + ivfflat）与共享 embeddings 客户端
  （app.core.llm_client.aembed）；嵌入维度(1024)、表结构、混合检索(RRF)全部沿用，不重复造轮子。
- 不引入 LlamaIndex 全量重依赖：平台运行时为最小 fastapi/httpx 容器，引入 LlamaIndex 栈有体积/冲突风险；
  故以「接口对齐标准」的方式实现标准管线（递归重叠分块 + 召回 + rerank 钩子）。后续若评估可接受，
  可平滑替换为 llama_index.core 的 VectorStoreIndex + PGVectorStore（对外接口一致）。

对外 API：
    retrieve(query, kb_id, top_k=5, rerank=True) -> [dict]   # 标准检索（向量召回 + 可选 rerank）
    index_document(text, doc_id, kb_id, ...) -> int          # 标准索引（递归重叠分块后逐块入库）
"""
import logging

from app.pgvector_search import vector_search, index_document as _pg_index

logger = logging.getLogger("aiflow.rag")

# 标准递归分块分隔符（与 LlamaIndex SentenceSplitter / LangChain RecursiveCharacterTextSplitter 同序）
_SEPARATORS = ["\n\n", "\n", "。", "！", "？", "；", ".", "!", "?", ";", " ", ""]


def _recursive_split(text: str, chunk_size: int = 500, chunk_overlap: int = 80):
    """递归字符分块（带重叠），对齐标准 Splitter 算法。

    逐层用分隔符切分；某段仍超长时递归细切；最终兜底硬切（带重叠步长）。
    """
    text = text or ""
    if len(text) <= chunk_size:
        return [text] if text.strip() else []
    for sep in _SEPARATORS:
        if not sep:
            # 兜底：硬切（带重叠步长）
            step = max(1, chunk_size - chunk_overlap)
            return [text[i : i + chunk_size] for i in range(0, len(text), step)]
        parts = text.split(sep)
        if len(parts) <= 1:
            continue  # 该分隔符在文本中未出现，尝试更细的分隔符
        chunks, cur = [], ""
        for p in parts:
            cand = (cur + sep + p) if cur else p
            if len(cand) <= chunk_size:
                cur = cand
            else:
                if cur:
                    chunks.append(cur)
                    cur = ""
                if len(p) > chunk_size:
                    chunks.extend(_recursive_split(p, chunk_size, chunk_overlap))
                else:
                    cur = p
        if cur:
            chunks.append(cur)
        if chunks:
            return [c for c in chunks if c.strip()]
    return [text]


async def index_document(
    text: str,
    doc_id: str,
    kb_id: str,
    chunk_size: int = 500,
    chunk_overlap: int = 80,
    metadata: dict | None = None,
) -> int:
    """标准索引：递归重叠分块后逐块存入 pgvector。返回成功分块数。"""
    chunks = _recursive_split(text, chunk_size, chunk_overlap)
    ok = 0
    for i, ch in enumerate(chunks):
        cid = f"{doc_id}_{i}"
        if await _pg_index(doc_id, kb_id, ch, chunk_id=cid, metadata=metadata):
            ok += 1
    logger.info(f"[rag] indexed {doc_id}: {ok}/{len(chunks)} chunks")
    return ok


async def retrieve(
    query: str,
    kb_id: str | None = None,
    top_k: int = 5,
    rerank: bool = True,
) -> list:
    """标准检索：向量召回 + 可选 rerank。

    返回 [{doc_id, chunk_id, content, score, metadata}, ...]。
    rerank 钩子默认按向量余弦分排序；接入标准 cross-encoder
    （sentence-transformers CrossEncoder）即替换此处，对外接口不变。
    """
    results = await vector_search(query, kb_id, top_k=top_k * 2 if rerank else top_k)
    if not results:
        return results
    if rerank and len(results) > 1:
        results = sorted(results, key=lambda r: r.get("score", 0), reverse=True)
    return results[:top_k]