"""U1 测试 — RAG 平台模块骨架 + LlamaIndex 集成。

测试场景：
1. LlamaIndex PGVectorStore 连接现有 pgvector 扩展（mock 验证参数）
2. 基础 ingest（文档 → chunk → embedding → pgvector INSERT）端到端工作
3. 基础 query（query → embedding → pgvector cosine 检索）返回结果
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

from agentkit.rag_platform.indexing import (
    DEFAULT_EMBED_DIM,
    KB_CHUNKS_TABLE,
    _async_to_sync_url,
    create_vector_store,
)
from agentkit.rag_platform.models import (
    Chunk,
    Document,
    DocumentStatus,
    KBStatus,
    KnowledgeBase,
    QueryMode,
    QueryResult,
)
from agentkit.rag_platform.pipeline import RAGPipeline


class TestModels:
    """领域模型测试。"""

    def test_knowledge_base_defaults(self):
        """KnowledgeBase 默认值正确。"""
        kb = KnowledgeBase(name="test-kb", owner="user1")
        assert kb.status == KBStatus.active
        assert kb.default_query_mode == QueryMode.blend
        assert kb.default_hit_processing == "model_opt"
        assert kb.caching_disabled is False
        assert kb.id  # auto-generated UUID

    def test_document_status_transitions(self):
        """DocumentStatus 状态机值正确。"""
        doc = Document(
            kb_id="kb1",
            filename="test.pdf",
            file_type="pdf",
            file_size=1024,
        )
        assert doc.status == DocumentStatus.pending
        assert doc.error_message is None

    def test_chunk_model(self):
        """Chunk 模型字段正确。"""
        chunk = Chunk(
            document_id="doc1",
            kb_id="kb1",
            content="hello world",
        )
        assert chunk.embedding is None
        assert chunk.metadata == {}

    def test_query_result_model(self):
        """QueryResult 模型字段正确。"""
        result = QueryResult(
            chunk_id="c1",
            content="hello",
            score=0.95,
            document_id="doc1",
            kb_id="kb1",
        )
        assert result.score == 0.95
        assert result.metadata == {}


class TestIndexing:
    """pgvector 索引管理测试。"""

    def test_async_to_sync_url_conversion(self):
        """asyncpg URL 正确转换为 psycopg2 URL。"""
        async_url = "postgresql+asyncpg://user:pass@localhost:5432/db"
        sync_url = _async_to_sync_url(async_url)
        assert sync_url == "postgresql://user:pass@localhost:5432/db"

    def test_sync_url_unchanged(self):
        """已经是 sync URL 时不转换。"""
        sync_url = "postgresql://user:pass@localhost:5432/db"
        assert _async_to_sync_url(sync_url) == sync_url

    @patch("llama_index.vector_stores.postgres.PGVectorStore.from_params")
    def test_create_vector_store_uses_explicit_table_name(self, mock_from_params):
        """create_vector_store 使用显式表名（schema 隔离）。"""
        mock_store = MagicMock()
        mock_from_params.return_value = mock_store

        result = create_vector_store(
            "postgresql+asyncpg://user:pass@localhost:5432/db",
            embed_dim=768,
        )

        # 验证 from_params 被调用时使用了正确的表名
        mock_from_params.assert_called_once()
        call_kwargs = mock_from_params.call_args.kwargs
        assert call_kwargs["table_name"] == KB_CHUNKS_TABLE
        assert call_kwargs["embed_dim"] == 768
        assert call_kwargs["hybrid_search"] is True
        assert result is mock_store

    @patch("llama_index.vector_stores.postgres.PGVectorStore.from_params")
    def test_create_vector_store_default_embed_dim(self, mock_from_params):
        """create_vector_store 默认 embed_dim 为 1536。"""
        mock_from_params.return_value = MagicMock()

        create_vector_store("postgresql://localhost/db")

        call_kwargs = mock_from_params.call_args.kwargs
        assert call_kwargs["embed_dim"] == DEFAULT_EMBED_DIM

    @patch("llama_index.vector_stores.postgres.PGVectorStore.from_params")
    def test_create_vector_store_converts_async_url(self, mock_from_params):
        """create_vector_store 将 asyncpg URL 转换为 sync URL。"""
        mock_from_params.return_value = MagicMock()

        create_vector_store("postgresql+asyncpg://user:pass@localhost:5432/db")

        call_kwargs = mock_from_params.call_args.kwargs
        assert call_kwargs["database"] == "postgresql://user:pass@localhost:5432/db"


class TestRAGPipeline:
    """RAGPipeline 管道测试。"""

    def _make_mock_embed_model(self):
        """创建 mock embedding 模型。"""
        mock = MagicMock()
        mock.aget_text_embedding = AsyncMock(return_value=[0.1] * 1536)
        return mock

    def _make_mock_vector_store(self):
        """创建 mock vector store。"""
        mock = MagicMock()
        mock.aquery = AsyncMock()
        return mock

    def _make_mock_text_node(self, node_id: str, content: str, metadata: dict | None = None):
        """创建 mock TextNode。"""
        node = MagicMock()
        node.node_id = node_id
        node.get_content.return_value = content
        node.metadata = metadata or {}
        return node

    async def test_ingest_calls_pipeline_arun(self):
        """ingest 调用 IngestionPipeline.arun 并返回 TextNode 列表。"""
        mock_vs = self._make_mock_vector_store()
        mock_embed = self._make_mock_embed_model()

        mock_nodes = [
            self._make_mock_text_node("n1", "chunk 1", {"kb_id": "kb1"}),
            self._make_mock_text_node("n2", "chunk 2", {"kb_id": "kb1"}),
        ]

        with patch("llama_index.core.ingestion.IngestionPipeline") as mock_pipeline_cls:
            mock_pipeline = MagicMock()
            mock_pipeline.arun = AsyncMock(return_value=mock_nodes)
            mock_pipeline_cls.return_value = mock_pipeline

            rag = RAGPipeline(mock_vs, mock_embed, chunk_size=256, chunk_overlap=20)
            nodes = await rag.ingest("hello world", metadata={"kb_id": "kb1", "document_id": "d1"})

        assert len(nodes) == 2
        assert nodes[0].node_id == "n1"
        mock_pipeline.arun.assert_awaited_once()

    async def test_query_returns_query_results(self):
        """query 返回 QueryResult 列表，按相似度排序。"""
        mock_vs = self._make_mock_vector_store()
        mock_embed = self._make_mock_embed_model()

        # 模拟 vector store query 返回
        mock_result = MagicMock()
        mock_result.nodes = [
            self._make_mock_text_node(
                "n1", "relevant chunk", {"kb_id": "kb1", "document_id": "d1"}
            ),
            self._make_mock_text_node("n2", "another chunk", {"kb_id": "kb1", "document_id": "d1"}),
        ]
        mock_result.similarities = [0.95, 0.80]
        mock_vs.aquery.return_value = mock_result

        with patch("llama_index.core.ingestion.IngestionPipeline"):
            rag = RAGPipeline(mock_vs, mock_embed)
            results = await rag.query("test query", top_k=2)

        assert len(results) == 2
        assert isinstance(results[0], QueryResult)
        assert results[0].chunk_id == "n1"
        assert results[0].score == 0.95
        assert results[0].kb_id == "kb1"
        assert results[1].score == 0.80

    async def test_query_calls_embed_model(self):
        """query 调用 embedding 模型获取查询向量。"""
        mock_vs = self._make_mock_vector_store()
        mock_embed = self._make_mock_embed_model()

        mock_result = MagicMock()
        mock_result.nodes = []
        mock_result.similarities = []
        mock_vs.aquery.return_value = mock_result

        with patch("llama_index.core.ingestion.IngestionPipeline"):
            rag = RAGPipeline(mock_vs, mock_embed)
            await rag.query("test", top_k=5)

        mock_embed.aget_text_embedding.assert_awaited_once_with("test")

    async def test_query_passes_top_k_to_vector_store(self):
        """query 将 top_k 传递给 vector store query。"""
        mock_vs = self._make_mock_vector_store()
        mock_embed = self._make_mock_embed_model()

        mock_result = MagicMock()
        mock_result.nodes = []
        mock_result.similarities = []
        mock_vs.aquery.return_value = mock_result

        with patch("llama_index.core.ingestion.IngestionPipeline"):
            rag = RAGPipeline(mock_vs, mock_embed)
            await rag.query("test", top_k=10)

        # 验证 aquery 被调用时 similarity_top_k=10
        call_args = mock_vs.aquery.call_args
        vs_query = call_args.args[0]
        assert vs_query.similarity_top_k == 10

    def test_pipeline_init_with_custom_chunk_params(self):
        """RAGPipeline 接受自定义 chunk_size 和 chunk_overlap。"""
        mock_vs = self._make_mock_vector_store()
        mock_embed = self._make_mock_embed_model()

        with patch("llama_index.core.ingestion.IngestionPipeline") as mock_pipeline_cls:
            RAGPipeline(mock_vs, mock_embed, chunk_size=1024, chunk_overlap=100)

        # 验证 SentenceSplitter 被创建时使用了自定义参数
        call_args = mock_pipeline_cls.call_args
        transformations = call_args.kwargs.get("transformations") or call_args[1].get(
            "transformations"
        )
        assert transformations is not None
        # 第一个 transformation 是 SentenceSplitter
        splitter = transformations[0]
        assert splitter.chunk_size == 1024
        assert splitter.chunk_overlap == 100
