"""知识的切片索引与检索（M3）。

两条硬规则：

1. **只有 `published` 的知识能被检索到**。草稿是"还没被人认可的话"，
   让它进检索等于把猜测当结论引用出去。过滤写在 SQL 里，不靠调用方自觉。
2. **检索必须带来源**。每个命中都返回 `item_id` 与标题，用户能点回去看原文；
   这和 D3（知识必须有来源）是同一条原则的延伸——不给"出处不存在"的答案。

评分是向量与词法的加权和：
- 向量负责"说的是同一件事但换了词"；
- 词法（pg_trgm 的 trigram 相似度）负责"专有名词、单号、人名"这类必须原样命中的词，
  而哈希向量对这类词并不敏感。
两者都不足时宁可返回空，也不硬凑一条不相关的结果——检索失败比误导便宜。
"""

from __future__ import annotations

import uuid
from typing import Any

from sqlalchemy import delete, select, text
from sqlalchemy.ext.asyncio import AsyncSession

from ..config import get_settings
from ..domain.chunking import split_markdown
from ..errors import AppError, ErrorCode
from ..models_knowledge import Chunk, KnowledgeDoc
from ..models_task import Item
from .embedding import get_provider

SEARCH_SQL = text(
    """
    SELECT c.item_id,
           i.title,
           k.doc_type,
           c.idx,
           c.heading,
           c.content,
           coalesce(app_cosine(c.embedding, CAST(:qvec AS real[])), 0) AS vec_score,
           coalesce(similarity(c.content, :q), 0) AS lex_score
    FROM chunk c
    JOIN item i ON i.id = c.item_id AND i.workspace_id = c.workspace_id
    JOIN knowledge_doc k ON k.item_id = c.item_id AND k.workspace_id = c.workspace_id
    WHERE c.workspace_id = :ws
      AND k.status_hint = 'published'
    ORDER BY (
        coalesce(app_cosine(c.embedding, CAST(:qvec AS real[])), 0) * :vw
        + coalesce(similarity(c.content, :q), 0) * :lw
    ) DESC, c.item_id, c.idx
    LIMIT :limit
    """
)


def _vector_param(values: list[float]) -> list[float]:
    """把向量交给驱动去编码。

    一开始我拼成 `'{0.1,0.2}'` 字符串再 `CAST(... AS real[])`，asyncpg 会因为 CAST
    已经告诉它"这是数组"而拒绝字符串（`a sized iterable container expected`）。
    直接传 list 最省事；非有限值替换成 0，免得 NaN 进 SQL 变成错误输入。
    """
    return [value if value == value and abs(value) != float("inf") else 0.0 for value in values]


async def reindex_document(
    session: AsyncSession,
    ws_id: uuid.UUID,
    item_id: uuid.UUID,
) -> int:
    """重建一份知识的切片，返回切片数。

    全量重建而不是增量：知识正文改动通常牵动上下几段，增量对齐的复杂度
    换不来多少收益（一份知识撑死几十片）。
    """
    item = (
        await session.execute(select(Item).where(Item.workspace_id == ws_id, Item.id == item_id))
    ).scalar_one_or_none()
    if item is None:
        raise AppError(ErrorCode.ITEM_NOT_FOUND, "这件事不存在")

    doc = (
        await session.execute(
            select(KnowledgeDoc).where(
                KnowledgeDoc.workspace_id == ws_id,
                KnowledgeDoc.item_id == item_id,
            )
        )
    ).scalar_one_or_none()
    if doc is None:
        raise AppError(ErrorCode.KNOWLEDGE_NOT_FOUND, "这份知识不存在")

    settings = get_settings()
    source = "\n\n".join(part for part in (item.title, item.body_md or "") if part)
    drafts = split_markdown(
        source,
        max_chars=settings.chunk_max_chars,
        overlap=settings.chunk_overlap_chars,
    )

    provider = get_provider(settings)
    vectors = await provider.embed([draft.content for draft in drafts]) if drafts else []

    await session.execute(
        delete(Chunk).where(Chunk.workspace_id == ws_id, Chunk.item_id == item_id)
    )
    for index, (draft, vector) in enumerate(zip(drafts, vectors, strict=True)):
        session.add(
            Chunk(
                workspace_id=ws_id,
                item_id=item_id,
                idx=index,
                heading=draft.heading,
                content=draft.content,
                char_count=len(draft.content),
                embedding=vector,
                embed_model=provider.name,
            )
        )
    doc.chunk_count = len(drafts)
    await session.flush()
    return len(drafts)


async def search_knowledge(
    session: AsyncSession,
    ws_id: uuid.UUID,
    query: str,
    *,
    limit: int = 10,
    mode: str = "hybrid",
) -> list[dict[str, Any]]:
    """检索已发布知识，返回按相关度排序的切片。

    `mode`: hybrid（默认）/ vector / keyword —— 前两者主要用于对比效果。
    """
    cleaned = (query or "").strip()
    if not cleaned:
        return []

    settings = get_settings()
    provider = get_provider(settings)
    query_vector = (await provider.embed([cleaned]))[0]

    if mode == "vector":
        vector_weight, lexical_weight = 1.0, 0.0
    elif mode == "keyword":
        vector_weight, lexical_weight = 0.0, 1.0
    else:
        vector_weight, lexical_weight = (
            settings.search_vector_weight,
            settings.search_lexical_weight,
        )

    rows = (
        await session.execute(
            SEARCH_SQL,
            {
                "qvec": _vector_param(query_vector),
                "q": cleaned,
                "ws": str(ws_id),
                "vw": vector_weight,
                "lw": lexical_weight,
                "limit": limit,
            },
        )
    ).mappings()

    hits: list[dict[str, Any]] = []
    for row in rows:
        vector_score = float(row["vec_score"] or 0.0)
        lexical_score = float(row["lex_score"] or 0.0)
        hits.append(
            {
                "item_id": row["item_id"],
                "title": row["title"],
                "doc_type": row["doc_type"],
                "idx": row["idx"],
                "heading": row["heading"],
                "content": row["content"],
                "score": round(vector_score * vector_weight + lexical_score * lexical_weight, 4),
                "vector_score": round(vector_score, 4),
                "lexical_score": round(lexical_score, 4),
            }
        )
    hits = [hit for hit in hits if hit["score"] >= settings.search_min_score]
    return hits


async def reindex_workspace(session: AsyncSession, ws_id: uuid.UUID) -> int:
    """把一个工作区里已发布的知识全部重建，返回处理的知识份数。

    换 embedding 模型之后必须整体跑一次，否则新旧向量混在一张表里，
    相似度分数没法比较。
    """
    item_ids = (
        await session.execute(
            select(KnowledgeDoc.item_id).where(
                KnowledgeDoc.workspace_id == ws_id,
                KnowledgeDoc.status_hint == "published",
            )
        )
    ).scalars()
    count = 0
    for item_id in item_ids:
        await reindex_document(session, ws_id, item_id)
        count += 1
    return count