"""LlamaIndex IngestionPipeline 封装 — 文档处理管道。

管道流程：文档 → SentenceSplitter 分段 → embedding → pgvector 索引。
U3 将扩展为完整 IngestionPipeline（含解析、预览、净化）。
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from agentkit.rag_platform.indexing import KB_CHUNKS_TABLE
from agentkit.rag_platform.models import QueryResult

if TYPE_CHECKING:
    from llama_index.core.embeddings import BaseEmbedding
    from llama_index.core.schema import TextNode
    from llama_index.vector_stores.postgres import PGVectorStore

logger = logging.getLogger(__name__)

DEFAULT_CHUNK_SIZE = 512
DEFAULT_CHUNK_OVERLAP = 50
DEFAULT_TOP_K = 5


class RAGPipeline:
    """封装 LlamaIndex IngestionPipeline 用于 KB 文档处理。

    Args:
        vector_store: LlamaIndex PGVectorStore 实例
        embed_model: LlamaIndex embedding 模型
        chunk_size: 分段大小（token 数）
        chunk_overlap: 分段重叠（token 数）
    """

    def __init__(
        self,
        vector_store: "PGVectorStore",
        embed_model: "BaseEmbedding",
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    ) -> None:
        from llama_index.core.ingestion import IngestionPipeline
        from llama_index.core.node_parser import SentenceSplitter

        self._vector_store = vector_store
        self._embed_model = embed_model

        self._pipeline = IngestionPipeline(
            transformations=[
                SentenceSplitter(
                    chunk_size=chunk_size,
                    chunk_overlap=chunk_overlap,
                ),
                embed_model,
            ],
            vector_store=vector_store,
        )
        logger.info(
            "RAGPipeline initialized: chunk_size=%d, chunk_overlap=%d, table=%s",
            chunk_size,
            chunk_overlap,
            KB_CHUNKS_TABLE,
        )

    async def ingest(self, text: str, metadata: dict[str, object] | None = None) -> list["TextNode"]:
        """将文本摄入向量存储。

        Args:
            text: 文档文本
            metadata: 附加到每个 chunk 的元数据（kb_id, document_id 等）

        Returns:
            LlamaIndex TextNode 列表（已写入 vector store）
        """
        from llama_index.core.schema import Document as LIDocument

        doc = LIDocument(text=text, metadata=metadata or {})
        nodes = await self._pipeline.arun(documents=[doc])
        logger.info("Ingested %d chunks from document", len(nodes))
        return nodes

    async def query(
        self,
        query_text: str,
        top_k: int = DEFAULT_TOP_K,
    ) -> list[QueryResult]:
        """语义检索。

        Args:
            query_text: 查询文本
            top_k: 返回结果数

        Returns:
            QueryResult 列表（按相似度降序）
        """
        from llama_index.core.vector_stores.types import VectorStoreQuery

        query_embedding = await self._embed_model.aget_text_embedding(query_text)

        vs_query = VectorStoreQuery(
            query_embedding=query_embedding,
            similarity_top_k=top_k,
        )
        result = await self._vector_store.aquery(vs_query)

        results: list[QueryResult] = []
        for node, score in zip(result.nodes, result.similarities):
            meta = node.metadata or {}
            results.append(
                QueryResult(
                    chunk_id=node.node_id,
                    content=node.get_content(),
                    score=float(score) if score is not None else 0.0,
                    metadata=meta,
                    document_id=meta.get("document_id", ""),
                    kb_id=meta.get("kb_id", ""),
                )
            )
        logger.info("Query returned %d results (top_k=%d)", len(results), top_k)
        return results
