"""U2 测试 — KB 持久化存储。

测试场景：
1. KB 元数据写入 PG，重启后仍存在（mock 验证 commit）
2. owner 用户可查询自己的 KB
3. delete_kb CASCADE 删除关联数据
4. 文档 CRUD 操作
"""

from __future__ import annotations

from contextlib import asynccontextmanager
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock

from agentkit.rag_platform.models import (
    DocumentStatus,
    KnowledgeBase,
)
from agentkit.rag_platform.store import KBStore


def _make_kb_row(**overrides):
    """创建模拟 KBModel 行。"""
    defaults = {
        "id": "kb-001",
        "name": "test-kb",
        "description": "",
        "owner": "user1",
        "status": "active",
        "default_query_mode": "blend",
        "default_hit_processing": "model_opt",
        "caching_disabled": False,
        "created_at": datetime.now(timezone.utc),
        "updated_at": datetime.now(timezone.utc),
    }
    defaults.update(overrides)
    mock = MagicMock()
    mock.configure_mock(**defaults)
    return mock


def _make_doc_row(**overrides):
    """创建模拟 DocumentModel 行。"""
    defaults = {
        "id": "doc-001",
        "kb_id": "kb-001",
        "filename": "test.pdf",
        "file_type": "pdf",
        "file_size": 1024,
        "status": "pending",
        "error_message": None,
        "created_at": datetime.now(timezone.utc),
        "updated_at": datetime.now(timezone.utc),
    }
    defaults.update(overrides)
    mock = MagicMock()
    mock.configure_mock(**defaults)
    return mock


def _make_mock_session_factory():
    """创建 mock session factory，返回 (factory, mock_session)。"""
    mock_session = AsyncMock()
    mock_session.add = MagicMock()
    mock_session.commit = AsyncMock()
    mock_session.rollback = AsyncMock()
    mock_session.refresh = AsyncMock()
    mock_session.delete = AsyncMock()
    mock_session.flush = AsyncMock()

    @asynccontextmanager
    async def factory():
        yield mock_session

    return factory, mock_session


def _make_mock_execute(result_rows):
    """创建 mock execute，返回 scalars().all() = result_rows。"""
    mock_result = MagicMock()
    mock_scalars = MagicMock()
    mock_scalars.all.return_value = result_rows
    mock_scalars.first.return_value = result_rows[0] if result_rows else None
    mock_result.scalars.return_value = mock_scalars
    return AsyncMock(return_value=mock_result)


class TestKBStoreCreate:
    """KB 创建测试。"""

    async def test_create_kb_returns_knowledge_base(self):
        """create_kb 返回 KnowledgeBase 领域模型。"""
        sf, mock_session = _make_mock_session_factory()

        # flush 后 kb.id 可用
        def _flush_side_effect():
            mock_session._kb_id = "kb-001"

        mock_session.flush.side_effect = _flush_side_effect

        # refresh 后设置所有必需属性
        async def _refresh(obj):
            obj.id = "kb-001"
            obj.status = "active"
            obj.created_at = datetime.now(timezone.utc)
            obj.updated_at = datetime.now(timezone.utc)

        mock_session.refresh.side_effect = _refresh

        store = KBStore(sf)
        kb = await store.create_kb(name="test", owner="user1", description="desc")

        assert isinstance(kb, KnowledgeBase)
        assert kb.name == "test"
        assert kb.owner == "user1"
        assert kb.description == "desc"
        # 验证 add 被调用（KB + ACL）
        assert mock_session.add.call_count == 2
        mock_session.commit.assert_awaited_once()

    async def test_create_kb_creates_owner_acl(self):
        """create_kb 在同一事务中创建 owner ACL 条目。"""
        sf, mock_session = _make_mock_session_factory()

        async def _refresh(obj):
            obj.id = "kb-001"
            obj.status = "active"
            obj.created_at = datetime.now(timezone.utc)
            obj.updated_at = datetime.now(timezone.utc)

        mock_session.refresh.side_effect = _refresh

        store = KBStore(sf)
        await store.create_kb(name="test", owner="user1")

        # 验证第二个 add 调用是 ACL 条目
        second_add_call = mock_session.add.call_args_list[1]
        acl_obj = second_add_call.args[0]
        assert acl_obj.user_id == "user1"
        assert acl_obj.role == "owner"


class TestKBStoreQuery:
    """KB 查询测试。"""

    async def test_get_kb_returns_none_if_not_found(self):
        """get_kb 返回 None 当 KB 不存在。"""
        sf, mock_session = _make_mock_session_factory()
        mock_session.execute = _make_mock_execute([])

        store = KBStore(sf)
        result = await store.get_kb("nonexistent")

        assert result is None

    async def test_get_kb_returns_knowledge_base(self):
        """get_kb 返回 KnowledgeBase 领域模型。"""
        sf, mock_session = _make_mock_session_factory()
        kb_row = _make_kb_row()
        mock_session.execute = _make_mock_execute([kb_row])

        store = KBStore(sf)
        result = await store.get_kb("kb-001")

        assert result is not None
        assert result.id == "kb-001"
        assert result.name == "test-kb"

    async def test_list_kbs_filters_by_owner(self):
        """list_kbs 按 owner 过滤。"""
        sf, mock_session = _make_mock_session_factory()
        kb_rows = [_make_kb_row(id="kb-1"), _make_kb_row(id="kb-2", name="second")]
        mock_session.execute = _make_mock_execute(kb_rows)

        store = KBStore(sf)
        results = await store.list_kbs(owner="user1")

        assert len(results) == 2
        assert results[0].id == "kb-1"

    async def test_list_kbs_returns_empty(self):
        """list_kbs 空列表正常返回。"""
        sf, mock_session = _make_mock_session_factory()
        mock_session.execute = _make_mock_execute([])

        store = KBStore(sf)
        results = await store.list_kbs()

        assert results == []


class TestKBStoreDelete:
    """KB 删除测试。"""

    async def test_delete_kb_returns_true(self):
        """delete_kb 删除成功返回 True。"""
        sf, mock_session = _make_mock_session_factory()
        kb_row = _make_kb_row()
        mock_session.execute = _make_mock_execute([kb_row])

        store = KBStore(sf)
        result = await store.delete_kb("kb-001")

        assert result is True
        mock_session.delete.assert_awaited_once_with(kb_row)
        mock_session.commit.assert_awaited_once()

    async def test_delete_kb_returns_false_if_not_found(self):
        """delete_kb KB 不存在返回 False。"""
        sf, mock_session = _make_mock_session_factory()
        mock_session.execute = _make_mock_execute([])

        store = KBStore(sf)
        result = await store.delete_kb("nonexistent")

        assert result is False


class TestDocumentOperations:
    """文档 CRUD 测试。"""

    async def test_add_document(self):
        """add_document 创建文档记录。"""
        sf, mock_session = _make_mock_session_factory()

        async def _refresh(obj):
            obj.id = "doc-001"
            obj.status = "pending"
            obj.created_at = datetime.now(timezone.utc)
            obj.updated_at = datetime.now(timezone.utc)

        mock_session.refresh.side_effect = _refresh

        store = KBStore(sf)
        doc = await store.add_document("kb-001", "test.pdf", "pdf", 1024)

        assert doc.filename == "test.pdf"
        assert doc.kb_id == "kb-001"
        mock_session.add.assert_called_once()
        mock_session.commit.assert_awaited_once()

    async def test_list_documents(self):
        """list_documents 返回指定 KB 的文档。"""
        sf, mock_session = _make_mock_session_factory()
        doc_rows = [_make_doc_row(id="d1"), _make_doc_row(id="d2", filename="second.pdf")]
        mock_session.execute = _make_mock_execute(doc_rows)

        store = KBStore(sf)
        results = await store.list_documents("kb-001")

        assert len(results) == 2
        assert results[0].id == "d1"

    async def test_update_document_status(self):
        """update_document_status 更新文档状态。"""
        sf, mock_session = _make_mock_session_factory()
        doc_row = _make_doc_row()
        mock_session.execute = _make_mock_execute([doc_row])

        store = KBStore(sf)
        result = await store.update_document_status("doc-001", DocumentStatus.indexed)

        assert result is not None
        assert doc_row.status == "indexed"
        mock_session.commit.assert_awaited_once()

    async def test_update_document_status_not_found(self):
        """update_document_status 文档不存在返回 None。"""
        sf, mock_session = _make_mock_session_factory()
        mock_session.execute = _make_mock_execute([])

        store = KBStore(sf)
        result = await store.update_document_status("nonexistent", DocumentStatus.failed)

        assert result is None

    async def test_delete_document(self):
        """delete_document 删除成功。"""
        sf, mock_session = _make_mock_session_factory()
        doc_row = _make_doc_row()
        mock_session.execute = _make_mock_execute([doc_row])

        store = KBStore(sf)
        result = await store.delete_document("doc-001")

        assert result is True
        mock_session.delete.assert_awaited_once_with(doc_row)
