"""文档处理管道 — 解析 → 分段 → 向量化。

封装 DocumentLoader（多格式解析）+ LlamaIndex SentenceSplitter（分段）+
PGVectorStore（向量化存储）。

状态机：pending → parsing → segmenting → vectorizing → indexed | failed
失败语义：vectorize 是 all-or-nothing — 失败时抛出异常，调用方负责将状态置为 failed。
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from agentkit.memory.document_loader import DocumentLoader
from agentkit.rag_platform.models import DocumentStatus
from agentkit.rag_platform.sanitize import sanitize_content

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

    from agentkit.rag_platform.store import KBStore

logger = logging.getLogger(__name__)

DEFAULT_CHUNK_SIZE = 512
DEFAULT_CHUNK_OVERLAP = 50


class DocumentProcessor:
    """文档处理管道：解析 → 分段 → 向量化。

    Args:
        chunk_size: 分段大小（token 数），默认 512
        chunk_overlap: 分段重叠（token 数），默认 50
    """

    def __init__(
        self,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    ) -> None:
        self._loader = DocumentLoader()
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap

    def parse(self, file_path: str, file_type: str) -> str:
        """从文件中提取文本。

        使用 memory/document_loader.py 的 DocumentLoader 进行多格式解析
        （PDF/DOCX/XLSX/Markdown/HTML/纯文本）。

        Args:
            file_path: 文件路径
            file_type: 文件类型（扩展名，不含点）— 用于净化

        Returns:
            提取的文本内容（已净化）

        Raises:
            FileNotFoundError: 文件不存在
            ValueError: 内容超过大小限制
        """
        doc = self._loader.load(file_path)
        # 对文本格式应用内容净化
        return sanitize_content(doc.content, file_type)

    def segment(
        self,
        text: str,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    ) -> list[str]:
        """将文本分段为 chunk。

        使用 LlamaIndex SentenceSplitter 进行分段。

        Args:
            text: 待分段文本
            chunk_size: 分段大小（token 数）
            chunk_overlap: 分段重叠（token 数）

        Returns:
            chunk 文本列表
        """
        from llama_index.core.node_parser import SentenceSplitter

        splitter = SentenceSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )
        chunks = splitter.split_text(text)
        logger.debug(
            "Segmented text into %d chunks (chunk_size=%d, overlap=%d)",
            len(chunks),
            chunk_size,
            chunk_overlap,
        )
        return chunks

    def preview(
        self,
        file_path: str,
        file_type: str,
        chunk_size: int | None = None,
        chunk_overlap: int | None = None,
    ) -> list[dict[str, object]]:
        """解析 + 分段，返回 chunk 用于只读预览（不向量化）。

        Args:
            file_path: 文件路径
            file_type: 文件类型（扩展名，不含点）
            chunk_size: 分段大小（默认使用构造器值）
            chunk_overlap: 分段重叠（默认使用构造器值）

        Returns:
            chunk 字典列表：[{index, content, metadata}, ...]
        """
        cs = chunk_size or self._chunk_size
        co = chunk_overlap or self._chunk_overlap

        text = self.parse(file_path, file_type)
        chunks = self.segment(text, cs, co)

        return [{"index": i, "content": c, "metadata": {}} for i, c in enumerate(chunks)]

    async def vectorize(
        self,
        chunks: list[str] | list[dict[str, object]],
        kb_id: str,
        document_id: str,
        vector_store: "PGVectorStore",
        embed_model: "BaseEmbedding",
    ) -> list["TextNode"]:
        """将 chunk 向量化并存储到 pgvector。

        All-or-nothing 语义：任一 chunk 向量化失败则整体失败（抛出异常）。

        Args:
            chunks: chunk 文本列表或 chunk 字典列表
            kb_id: 知识库 ID
            document_id: 文档 ID
            vector_store: LlamaIndex PGVectorStore 实例
            embed_model: LlamaIndex embedding 模型

        Returns:
            已写入 vector store 的 TextNode 列表

        Raises:
            Exception: 向量化或存储失败
        """
        from llama_index.core.schema import TextNode

        # 统一 chunk 格式为字符串
        chunk_texts: list[str] = []
        for i, chunk in enumerate(chunks):
            if isinstance(chunk, str):
                chunk_texts.append(chunk)
            elif isinstance(chunk, dict):
                chunk_texts.append(chunk.get("content", ""))
            else:
                chunk_texts.append(str(chunk))

        # 构建 TextNode
        nodes: list[TextNode] = [
            TextNode(
                text=text,
                metadata={
                    "kb_id": kb_id,
                    "document_id": document_id,
                    "chunk_index": i,
                },
            )
            for i, text in enumerate(chunk_texts)
        ]

        # 向量化 — all-or-nothing：任一失败则抛出异常
        for node in nodes:
            node.embedding = await embed_model.aget_text_embedding(node.text)

        # 存储到 vector store
        await vector_store.async_add(nodes)

        logger.info(
            "Vectorized %d chunks for document=%s kb=%s",
            len(nodes),
            document_id,
            kb_id,
        )
        return nodes

    async def process(
        self,
        file_path: str,
        file_type: str,
        kb_id: str,
        document_id: str,
        vector_store: "PGVectorStore",
        embed_model: "BaseEmbedding",
        store: "KBStore",
    ) -> None:
        """完整管道：parse → segment → vectorize，含状态转换。

        状态转换：pending → parsing → segmenting → vectorizing → indexed
        失败时：状态置为 failed，error_message 记录异常信息，异常重新抛出。

        Args:
            file_path: 文件路径
            file_type: 文件类型（扩展名，不含点）
            kb_id: 知识库 ID
            document_id: 文档 ID
            vector_store: LlamaIndex PGVectorStore 实例
            embed_model: LlamaIndex embedding 模型
            store: KBStore 实例（用于状态更新）

        Raises:
            Exception: 管道任一阶段失败（状态已置为 failed）
        """
        try:
            # parsing
            await store.update_document_status(document_id, DocumentStatus.parsing)
            text = self.parse(file_path, file_type)

            # segmenting
            await store.update_document_status(document_id, DocumentStatus.segmenting)
            chunks = self.segment(text, self._chunk_size, self._chunk_overlap)

            # vectorizing
            await store.update_document_status(document_id, DocumentStatus.vectorizing)
            await self.vectorize(chunks, kb_id, document_id, vector_store, embed_model)

            # indexed
            await store.update_document_status(document_id, DocumentStatus.indexed)
            logger.info("Document %s processed successfully", document_id)
        except Exception as e:
            # failed — 记录错误信息并重新抛出
            await store.update_document_status(
                document_id, DocumentStatus.failed, error_message=str(e)
            )
            logger.error("Document %s processing failed: %s", document_id, e)
            raise


__all__ = [
    "DEFAULT_CHUNK_OVERLAP",
    "DEFAULT_CHUNK_SIZE",
    "DocumentProcessor",
]
