"""PG 全文检索 — jieba 分词 + tsvector 写入/查询。

避免依赖 PG 扩展（pg_jieba/zhparser）：在 Python 层用 jieba 分词后，
将 token 用空格连接写入 `search_vector` 列（`to_tsvector('simple', ...)`）。
查询时同样用 jieba 分词构造 tsquery（AND 语义）。

依赖 `rag_platform_kb_chunks` 表的 `search_vector` 列（由 PGVectorStore
hybrid_search=True 自动创建）。
"""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING

from sqlalchemy import text

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# 表名 — 与 indexing.py 保持一致
KB_CHUNKS_TABLE = "rag_platform_kb_chunks"

# ponytail: tsquery 中需转义的字符（PG to_tsquery 语法）
# 升级路径：若需支持短语查询/权重，改用 phraseto_tsquery
_TSQUERY_SPECIAL = re.compile(r"[&|!():<>\"'\\]")


def tokenize(text: str) -> str:
    """jieba 分词后用空格连接 — 用于 tsvector 写入。

    精确模式（cut_all=False）适合索引构建。
    """
    import jieba

    tokens = jieba.cut(text, cut_all=False)
    return " ".join(tokens)


def build_tsquery(text: str) -> str:
    """jieba 分词后构造 tsquery — 用于全文检索查询。

    用 `&` 连接（AND 语义），过滤空 token 和纯空白 token。
    转义 PG to_tsquery 的特殊字符，避免语法错误。
    """
    import jieba

    tokens = jieba.cut(text, cut_all=False)
    cleaned: list[str] = []
    for t in tokens:
        t = t.strip()
        if not t:
            continue
        # 转义 PG to_tsquery 特殊字符
        t = _TSQUERY_SPECIAL.sub(" ", t).strip()
        if t:
            cleaned.append(t)
    return " & ".join(cleaned)


async def write_search_vector(session: "AsyncSession", chunk_id: str, content: str) -> None:
    """将 jieba 分词后的内容写入 chunk 的 search_vector 列。

    使用 `to_tsvector('simple', space_joined_tokens)` 写入 — 'simple' 配置
    不做语干提取，保留 jieba 切分的原始 token。

    Args:
        session: SQLAlchemy async session（调用方负责 commit）
        chunk_id: chunk 行 ID（PGVectorStore 写入后的 node_id）
        content: chunk 原始文本
    """
    tokenized = tokenize(content)
    await session.execute(
        text(
            f"UPDATE {KB_CHUNKS_TABLE} "  # noqa: S608 — 表名为常量，无注入风险
            "SET search_vector = to_tsvector('simple', :tokens) "
            "WHERE id = :chunk_id"
        ),
        {"tokens": tokenized, "chunk_id": chunk_id},
    )


async def write_search_vector_batch(session: "AsyncSession", items: list[tuple[str, str]]) -> None:
    """批量写入 search_vector — 用于文档索引构建后批量回填。

    Args:
        session: SQLAlchemy async session（调用方负责 commit）
        items: [(chunk_id, content), ...]
    """
    for chunk_id, content in items:
        await write_search_vector(session, chunk_id, content)


__all__ = [
    "build_tsquery",
    "tokenize",
    "write_search_vector",
    "write_search_vector_batch",
]
