"""
TreeSearch — 结构化文档两级检索
================================

为知识库提供有结构的「两级检索」能力（零额外基础设施成本，复用 MySQL 已存的分段数据）：

  Level 1 (build_kb_tree)   : 知识库 -> 目录 -> 文档 的嵌套结构树，用于浏览/定位/面包屑。
  Level 2 (tree_search)     : 在给定作用域（整个 KB / 某目录 / 某文档）内做分段(chunk)级检索，
                               BM25 关键词排序 + 可选 pgvector 语义向量 boost，返回带完整
                               面包屑(KB名 -> 目录路径 -> 文件名 -> 段序号/标题)的精准命中。

设计要点：
  * Level 2 直接消费 knowledge_document_chunks 表（content + chunk_metadata.title/keywords），
    因此召回的是「文档内部的精确片段」，而非整篇文档——这正是 RAG 想要的细粒度上下文。
  * 与 kb_router_search 融合：tree_search 作为 RRF 的一个新来源 (source="tree_chunk") 接入，
    聊天 KB 上下文自动获得结构化分片段命中。
  * 向量 boost 为可选且优雅降级：当 pgvector 无对应数据时，自动退化为纯 BM25，不影响主链路。
"""

import logging
import math
from typing import List, Dict, Optional

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


# ---------------------------------------------------------------------------
# 中文友好的分词（jieba 优先，否则按字符 bigram + 英文正则兜底）
# ---------------------------------------------------------------------------
def _tokenize(text: str) -> List[str]:
    if not text:
        return []
    toks: List[str] = []
    try:
        import re
        # 英文/数字 token
        for m in re.findall(r"[a-zA-Z0-9_]+", text.lower()):
            if len(m) >= 1:
                toks.append(m)
        # 中文按 jieba 或字符切分
        cjk = "".join(re.findall(r"[一-鿿]", text))
        if cjk:
            try:
                import jieba
                for w in jieba.lcut(cjk):
                    w = w.strip().lower()
                    if w:
                        toks.append(w)
            except Exception:
                # 字符级 bigram（兼顾召回与精度）
                for i in range(len(cjk)):
                    toks.append(cjk[i])
                    if i + 1 < len(cjk):
                        toks.append(cjk[i:i + 2])
    except Exception:
        toks = [t for t in text.lower().split() if t]
    return toks


def _chunk_display_text(chunk) -> str:
    """拼装分段的检索文本：正文 + 标题 + 关键词（来自 chunk_metadata）。"""
    parts = []
    content = getattr(chunk, "content", None) or ""
    parts.append(content)
    md = getattr(chunk, "chunk_metadata", None) or {}
    if isinstance(md, dict):
        title = md.get("title") or md.get("heading") or ""
        keywords = md.get("keywords") or []
        if title:
            parts.append(title)
        if isinstance(keywords, list) and keywords:
            parts.append(" ".join(str(k) for k in keywords))
        elif isinstance(keywords, str) and keywords:
            parts.append(keywords)
    return "\n".join(p for p in parts if p).strip()


def _chunk_title(chunk) -> str:
    md = getattr(chunk, "chunk_metadata", None) or {}
    if isinstance(md, dict):
        return md.get("title") or md.get("heading") or ""
    return ""


# ---------------------------------------------------------------------------
# BM25（在作用域内的分段语料上即时计算 IDF）
# ---------------------------------------------------------------------------
def _bm25_rank(query_tokens: List[str], corpus: List[Dict], k1: float = 1.5, b: float = 0.75) -> List[Dict]:
    """corpus: [{"id":..., "tokens":[...], "chunk":<orm>, "text":str}] -> 返回带 bm25 分的列表"""
    if not query_tokens or not corpus:
        return []
    n = len(corpus)
    df = {}
    for qt in set(query_tokens):
        df[qt] = sum(1 for d in corpus if qt in d["tokens"])
    idf = {qt: math.log((n - df[qt] + 0.5) / (df[qt] + 0.5) + 1) for qt in query_tokens}
    avgdl = sum(len(d["tokens"]) for d in corpus) / max(n, 1)
    scored = []
    for d in corpus:
        dl = len(d["tokens"])
        score = 0.0
        for qt in set(query_tokens):
            if qt in d["tokens"]:
                f = d["tokens"].count(qt)
                score += idf[qt] * (f * (k1 + 1)) / (f + k1 * (1 - b + b * (dl / max(avgdl, 1))))
        if score > 0:
            scored.append({**d, "bm25": score})
    scored.sort(key=lambda x: x["bm25"], reverse=True)
    return scored


# ---------------------------------------------------------------------------
# Level 1：知识库目录树
# ---------------------------------------------------------------------------
def build_kb_tree(db, user_id: Optional[str] = None, kb_id: Optional[str] = None) -> List[Dict]:
    """构造 知识库 -> 目录(可嵌套) -> 文档 的结构树。"""
    from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument

    q = db.query(KnowledgeBase)
    if user_id:
        q = q.filter(KnowledgeBase.user_id == user_id)
    if kb_id:
        q = q.filter(KnowledgeBase.id == kb_id)
    bases = q.all()

    tree = []
    for kb in bases:
        dirs = db.query(KnowledgeDirectory).filter(
            KnowledgeDirectory.knowledge_base_id == kb.id
        ).all()
        docs = db.query(KnowledgeDocument).filter(
            KnowledgeDocument.knowledge_base_id == kb.id,
            KnowledgeDocument.status == "completed",
        ).all()
        dir_by_id = {d.id: d for d in dirs}
        node_by_id = {}

        def make_dir_node(d):
            return {
                "id": d.id,
                "name": d.name,
                "type": "directory",
                "children": [],
                "documents": [],
            }

        for d in dirs:
            node_by_id[d.id] = make_dir_node(d)
        root_dirs = []
        for d in dirs:
            node = node_by_id[d.id]
            parent = d.parent_id
            if parent and parent in node_by_id:
                node_by_id[parent]["children"].append(node)
            else:
                root_dirs.append(node)
        for doc in docs:
            item = {
                "id": doc.id,
                "name": doc.filename,
                "type": "document",
                "file_type": doc.file_type,
                "chunk_count": doc.chunk_count or 0,
            }
            if doc.directory_id and doc.directory_id in node_by_id:
                node_by_id[doc.directory_id]["documents"].append(item)
            else:
                root_dirs.append(item)  # 顶层文档
        tree.append({
            "id": kb.id,
            "name": kb.name,
            "type": "knowledge_base",
            "directories": root_dirs,
            "document_count": len(docs),
        })
    return tree


# ---------------------------------------------------------------------------
# 面包屑：KB名 -> 目录路径 -> 文件名 -> 段标题/序号
# ---------------------------------------------------------------------------
def _breadcrumb(db, kb_id: str, directory_id: Optional[str], doc_id: str, chunk_index: int, title: str) -> Dict:
    from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument
    kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    kb_name = kb.name if kb else kb_id
    path = []
    cur = directory_id
    guard = 0
    while cur and guard < 20:
        d = db.query(KnowledgeDirectory).filter(KnowledgeDirectory.id == cur).first()
        if not d:
            break
        path.insert(0, d.name)
        cur = d.parent_id
        guard += 1
    doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
    doc_name = doc.filename if doc else doc_id
    return {
        "knowledge_base": kb_name,
        "directory_path": "/".join(path) if path else "/",
        "document": doc_name,
        "chunk_index": chunk_index,
        "chunk_title": title,
        "label": f"{kb_name}/{('/'.join(path) if path else '')}/{doc_name}" + (f"#段{chunk_index}" + (f"·{title}" if title else "")),
    }


# ---------------------------------------------------------------------------
# Level 2：结构化分片段检索
# ---------------------------------------------------------------------------
async def tree_search(
    db,
    query: str,
    user_id: Optional[str] = None,
    kb_ids: Optional[List[str]] = None,
    scope_base: Optional[str] = None,
    scope_dir: Optional[str] = None,
    scope_doc: Optional[str] = None,
    top_k: int = 5,
    mode: str = "hybrid",
    max_tokens: int = 1200,
) -> List[Dict]:
    """
    两级检索的 Level 2：在作用域内做分片段精准检索。

    Args:
        query      : 查询文本
        kb_ids     : 限定知识库范围（多个）
        scope_base : 限定单一知识库
        scope_dir  : 限定单一目录（含子目录）
        scope_doc  : 限定单一文档
        top_k      : 返回命中数
        mode       : "bm25" | "vector" | "hybrid"(默认)
    Returns:
        [{
            "content": <分段文本>,
            "score": 0-1,
            "source": "tree_chunk",
            "metadata": { "breadcrumb": {...}, "snippet": ..., "chunk_index": int, "title": str }
        }, ...]
    """
    from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument, KnowledgeDocumentChunk

    if not query or not query.strip():
        return []

    # 1) 解析作用域 -> 目标文档 id 集合
    target_kb_ids = list(kb_ids or [])
    if scope_base and scope_base not in target_kb_ids:
        target_kb_ids.append(scope_base)

    doc_scope = set()
    if scope_doc:
        doc_scope.add(scope_doc)
    if scope_dir:
        # 收集该目录及其所有子目录下的文档
        dirs = db.query(KnowledgeDirectory).all()
        dir_by_id = {d.id: d for d in dirs}
        sub = set()

        def collect(did):
            sub.add(did)
            for d in dirs:
                if d.parent_id == did:
                    collect(d.id)
        collect(scope_dir)
        docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.directory_id.in_(sub)).all()
        for d in docs:
            doc_scope.add(d.id)
    if not doc_scope and not target_kb_ids and not user_id:
        return []

    # 2) 取分段语料
    q = db.query(KnowledgeDocumentChunk)
    if doc_scope:
        q = q.filter(KnowledgeDocumentChunk.document_id.in_(doc_scope))
    else:
        if target_kb_ids:
            q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(target_kb_ids))
        elif user_id:
            kb_ids_u = [r[0] for r in db.query(KnowledgeBase.id).filter(KnowledgeBase.user_id == user_id).all()]
            if not kb_ids_u:
                return []
            q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(kb_ids_u))
    chunks = q.order_by(KnowledgeDocumentChunk.chunk_index).limit(500).all()
    if not chunks:
        return []

    query_tokens = _tokenize(query)
    corpus = []
    for c in chunks:
        text = _chunk_display_text(c)
        if len(text.strip()) < 8:
            continue
        corpus.append({
            "id": c.id,
            "tokens": _tokenize(text),
            "chunk": c,
            "text": text,
        })

    # 3) BM25 排序
    bm25_hits = _bm25_rank(query_tokens, corpus) if mode in ("bm25", "hybrid") else []

    # 4) 可选向量 boost
    vector_boost = {}
    if mode in ("vector", "hybrid"):
        try:
            from app.pgvector_search import vector_search
            for kid in (target_kb_ids or [None]):
                try:
                    vres = await vector_search(query, kid, top_k=top_k * 3)
                except Exception as e:
                    logger.warning(f"[tree_search] vector_search failed: {e}")
                    continue
                for vr in vres:
                    vcontent = (vr.get("content") or "").strip()
                    if vcontent:
                        vector_boost[vcontent] = max(vector_boost.get(vcontent, 0.0), float(vr.get("score", 0.0)))
        except Exception as e:
            logger.debug(f"[tree_search] vector unavailable: {e}")

    # 5) 融合
    merged_map = {}
    for rank, h in enumerate(bm25_hits):
        key = h["text"][:300]
        boost = 0.0
        for vc, vs in vector_boost.items():
            if vc and vc in h["text"]:
                boost = max(boost, vs)
        merged_map[key] = {
            "chunk": h["chunk"],
            "bm25": h["bm25"],
            "vec": boost,
            "rank": rank + 1,
        }
    # 纯向量模式：把仅在向量里命中的也纳入
    if mode == "vector":
        for vc, vs in vector_boost.items():
            if vc[:300] not in merged_map:
                # 找到对应 chunk
                for c in corpus:
                    if vc in c["text"]:
                        merged_map[vc[:300]] = {"chunk": c["chunk"], "bm25": 0.0, "vec": vs, "rank": 999}
                        break

    if not merged_map:
        return []

    # 综合分数（BM25 主导，向量做 0~1 归一 boost）
    max_bm25 = max((m["bm25"] for m in merged_map.values()), default=1.0) or 1.0
    items = []
    for key, m in merged_map.items():
        c = m["chunk"]
        norm_bm25 = (m["bm25"] / max_bm25) if max_bm25 > 0 else 0.0
        norm_vec = min(m["vec"], 1.0)
        combined = norm_bm25 * 0.7 + norm_vec * 0.3
        title = _chunk_title(c)
        breadcrumb = _breadcrumb(db, c.knowledge_base_id, c.document_id and _doc_dir(db, c.document_id), c.document_id, c.chunk_index, title)
        snippet = _make_snippet(c.content or "", query)
        items.append({
            "content": c.content or "",
            "score": round(combined, 4),
            "source": "tree_chunk",
            "metadata": {
                "breadcrumb": breadcrumb,
                "snippet": snippet,
                "chunk_index": c.chunk_index,
                "title": title,
                "document_id": c.document_id,
                "knowledge_base_id": c.knowledge_base_id,
            },
        })
    items.sort(key=lambda x: x["score"], reverse=True)
    return items[:top_k]


def _doc_dir(db, doc_id: str):
    from app.models import KnowledgeDocument
    d = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
    return d.directory_id if d else None


def _make_snippet(content: str, query: str, win: int = 80) -> str:
    if not content:
        return ""
    q = query.strip()
    idx = content.find(q)
    if idx < 0:
        # 尝试首个 token
        toks = _tokenize(q)
        for t in toks:
            idx = content.find(t)
            if idx >= 0:
                break
    if idx < 0:
        return content[:win * 2]
    start = max(0, idx - win)
    end = min(len(content), idx + len(q) + win)
    snip = content[start:end]
    if start > 0:
        snip = "..." + snip
    if end < len(content):
        snip = snip + "..."
    return snip