"""U4 测试: 记忆系统 - 三层记忆 + 混合检索 + BaseAgent 生命周期集成"""

import math
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock

import pytest

from agentkit.core.base import BaseAgent
from agentkit.core.protocol import AgentCapability, TaskMessage, TaskResult, TaskStatus
from agentkit.memory.base import Memory, MemoryItem
from agentkit.memory.episodic import EpisodicMemory
from agentkit.memory.retriever import MemoryRetriever
from agentkit.memory.semantic import SemanticMemory
from agentkit.memory.working import WorkingMemory


# ── In-Memory Memory 实现（用于测试） ────────────────────


class InMemoryMemory(Memory):
    """基于内存的 Memory 实现，用于测试"""

    def __init__(self):
        self._store: dict[str, MemoryItem] = {}

    async def store(self, key: str, value, metadata=None) -> None:
        self._store[key] = MemoryItem(
            key=key, value=value, metadata=metadata or {}, score=1.0
        )

    async def retrieve(self, key: str) -> MemoryItem | None:
        return self._store.get(key)

    async def search(self, query: str, top_k: int = 5, filters=None) -> list[MemoryItem]:
        results = []
        for item in self._store.values():
            if query.lower() in str(item.value).lower() or query.lower() in item.key.lower():
                results.append(item)
        return results[:top_k]

    async def delete(self, key: str) -> bool:
        return self._store.pop(key, None) is not None


# ── Memory 基类测试 ──────────────────────────────────────


class TestMemoryBase:
    async def test_in_memory_store_and_retrieve(self):
        mem = InMemoryMemory()
        await mem.store("key1", {"data": "hello"}, {"tag": "test"})
        item = await mem.retrieve("key1")
        assert item is not None
        assert item.value["data"] == "hello"
        assert item.metadata["tag"] == "test"

    async def test_in_memory_search(self):
        mem = InMemoryMemory()
        await mem.store("agent:task1", "Generated content about AI")
        await mem.store("agent:task2", "Analysis of trends")
        await mem.store("agent:task3", "AI research summary")

        results = await mem.search("AI")
        assert len(results) == 2

    async def test_in_memory_delete(self):
        mem = InMemoryMemory()
        await mem.store("key1", "value1")
        assert await mem.delete("key1") is True
        assert await mem.retrieve("key1") is None
        assert await mem.delete("nonexistent") is False

    async def test_batch_store(self):
        mem = InMemoryMemory()
        items = [("k1", "v1", None), ("k2", "v2", {"tag": "t"})]
        await mem.store_batch(items)
        assert await mem.retrieve("k1") is not None
        assert await mem.retrieve("k2") is not None

    async def test_get_context(self):
        mem = InMemoryMemory()
        await mem.store("k1", "First context item about Python")
        await mem.store("k2", "Second context item about AI")
        context = await mem.get_context("Python")
        assert "Python" in context


# ── SemanticMemory 测试 ──────────────────────────────────


class TestSemanticMemory:
    async def test_rag_search(self):
        """通过 RAG 服务检索知识"""

        class MockRAGService:
            async def search(self, query, knowledge_base_ids=None, top_k=5):
                return [
                    {"id": "doc1", "content": f"Knowledge about {query}", "score": 0.9, "source": "kb1"},
                ]

        mem = SemanticMemory(rag_service=MockRAGService(), knowledge_base_ids=["kb1"])
        results = await mem.search("AI trends")
        assert len(results) == 1
        assert results[0].value == "Knowledge about AI trends"
        assert results[0].score == 0.9

    async def test_graph_search(self):
        """通过知识图谱检索"""

        class MockGraphService:
            async def query(self, query, depth=2):
                return [
                    {"id": "node1", "content": f"Entity: {query}", "score": 0.7, "entities": ["AI"]},
                ]

        mem = SemanticMemory(graph_service=MockGraphService())
        results = await mem.search("machine learning")
        assert len(results) == 1
        assert results[0].metadata["source"] == "graph"

    async def test_combined_rag_and_graph(self):
        """RAG + 图谱联合检索"""

        class MockRAG:
            async def search(self, query, **kwargs):
                return [{"id": "r1", "content": "RAG result", "score": 0.8, "source": "rag"}]

        class MockGraph:
            async def query(self, query, **kwargs):
                return [{"id": "g1", "content": "Graph result", "score": 0.6, "source": "graph"}]

        mem = SemanticMemory(rag_service=MockRAG(), graph_service=MockGraph())
        results = await mem.search("test")
        assert len(results) == 2
        # 按 score 排序
        assert results[0].score >= results[1].score

    async def test_semantic_read_only(self):
        """Semantic Memory 通常只读"""
        mem = SemanticMemory()
        assert await mem.delete("any") is False


# ── EpisodicMemory 测试（使用 mock ORM） ─────────────────


class TestEpisodicMemory:
    async def test_time_decay(self):
        """时间衰减：近期经验权重高于远期"""
        # 直接测试衰减公式
        decay_rate = 0.01
        now = datetime.now(timezone.utc)

        recent_score = 0.8 * math.exp(-decay_rate * 1)  # 1 hour ago
        old_score = 0.8 * math.exp(-decay_rate * 100)  # 100 hours ago

        assert recent_score > old_score
        assert recent_score > 0.7
        assert old_score < 0.5


# ── MemoryRetriever 测试 ─────────────────────────────────


class TestMemoryRetriever:
    async def test_retriever_with_in_memory(self):
        """混合检索器使用 InMemoryMemory"""
        working = InMemoryMemory()
        await working.store("current_task", "Working on AI content generation")

        retriever = MemoryRetriever(working_memory=working)
        results = await retriever.retrieve("AI content")
        assert len(results) >= 1

    async def test_retriever_weights(self):
        """不同层权重影响排序"""
        working = InMemoryMemory()
        semantic = InMemoryMemory()

        await working.store("task1", "Working memory result")
        await semantic.store("doc1", "Semantic memory result")

        retriever = MemoryRetriever(
            working_memory=working,
            semantic_memory=semantic,
            weights={"working": 0.2, "semantic": 0.8},
        )
        results = await retriever.retrieve("result")
        # Semantic 权重更高，应排前面
        if len(results) >= 2:
            assert results[0].score >= results[1].score

    async def test_retriever_token_budget(self):
        """Token 预算管理"""
        working = InMemoryMemory()
        for i in range(20):
            await working.store(f"item_{i}", f"Long content item number {i} " * 50)

        retriever = MemoryRetriever(working_memory=working)
        results = await retriever.retrieve("content", token_budget=200)
        total_chars = sum(len(str(r.value)) for r in results)
        # 粗略估算 token 数不应远超预算
        assert total_chars // 4 <= 250  # 允许少量溢出

    async def test_get_context_string(self):
        """获取格式化上下文字符串"""
        working = InMemoryMemory()
        await working.store("ctx1", "Context about Python programming")

        retriever = MemoryRetriever(working_memory=working)
        context = await retriever.get_context_string("Python")
        assert "Python" in context

    async def test_empty_retriever(self):
        """无记忆层时不报错"""
        retriever = MemoryRetriever()
        results = await retriever.retrieve("anything")
        assert results == []


# ── BaseAgent + Memory 生命周期集成测试 ───────────────────


class TestAgentMemoryIntegration:
    async def test_memory_injected_into_agent(self):
        """Memory 可注入到 Agent"""

        class TestAgent(BaseAgent):
            async def handle_task(self, task):
                return {"result": "done"}

            def get_capabilities(self):
                return AgentCapability(
                    agent_name=self.name, agent_type="test", version="1.0",
                    supported_tasks=["test"], max_concurrency=1, description="test",
                )

        agent = TestAgent(name="test", agent_type="test")
        mem = InMemoryMemory()
        agent.use_memory(mem)
        assert agent.memory is mem

    async def test_on_task_complete_stores_memory(self):
        """on_task_complete 钩子可存储记忆"""

        class MemoryAwareAgent(BaseAgent):
            async def handle_task(self, task):
                return {"answer": "42"}

            def get_capabilities(self):
                return AgentCapability(
                    agent_name=self.name, agent_type="test", version="1.0",
                    supported_tasks=["test"], max_concurrency=1, description="test",
                )

            async def on_task_complete(self, task, output):
                if self.memory:
                    await self.memory.store(
                        f"task:{task.task_id}",
                        output,
                        {"agent_name": self.name, "task_type": task.task_type},
                    )

        mem = InMemoryMemory()
        agent = MemoryAwareAgent(name="mem_agent", agent_type="test")
        agent.use_memory(mem)

        task = TaskMessage(
            task_id="t-001", agent_name="mem_agent", task_type="test",
            priority=1, input_data={}, callback_url=None,
            created_at=datetime.now(timezone.utc),
        )
        result = await agent.execute(task)
        assert result.status == TaskStatus.COMPLETED

        # 验证记忆已存储
        stored = await mem.retrieve("task:t-001")
        assert stored is not None
        assert stored.value["answer"] == "42"

    async def test_on_task_start_loads_memory(self):
        """on_task_start 钩子可加载记忆到任务上下文"""

        class ContextAwareAgent(BaseAgent):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.loaded_context = None

            async def handle_task(self, task):
                return {"context_used": self.loaded_context is not None}

            def get_capabilities(self):
                return AgentCapability(
                    agent_name=self.name, agent_type="test", version="1.0",
                    supported_tasks=["test"], max_concurrency=1, description="test",
                )

            async def on_task_start(self, task):
                if self.memory:
                    context = await self.memory.get_context(task.task_type)
                    self.loaded_context = context

        mem = InMemoryMemory()
        await mem.store("test", "Previous experience with similar tasks")

        agent = ContextAwareAgent(name="ctx_agent", agent_type="test")
        agent.use_memory(mem)

        task = TaskMessage(
            task_id="t-002", agent_name="ctx_agent", task_type="test",
            priority=1, input_data={}, callback_url=None,
            created_at=datetime.now(timezone.utc),
        )
        result = await agent.execute(task)
        assert result.output_data["context_used"] is True

    async def test_on_task_failed_records_failure(self):
        """on_task_failed 钩子可记录失败模式"""

        class ResilientAgent(BaseAgent):
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.failure_recorded = False

            async def handle_task(self, task):
                raise ValueError("simulated failure")

            def get_capabilities(self):
                return AgentCapability(
                    agent_name=self.name, agent_type="test", version="1.0",
                    supported_tasks=["test"], max_concurrency=1, description="test",
                )

            async def on_task_failed(self, task, error):
                if self.memory:
                    await self.memory.store(
                        f"failure:{task.task_id}",
                        {"error": str(error), "task_type": task.task_type},
                        {"outcome": "failure"},
                    )
                self.failure_recorded = True

        mem = InMemoryMemory()
        agent = ResilientAgent(name="resilient", agent_type="test")
        agent.use_memory(mem)

        task = TaskMessage(
            task_id="t-003", agent_name="resilient", task_type="test",
            priority=1, input_data={}, callback_url=None,
            created_at=datetime.now(timezone.utc),
        )
        result = await agent.execute(task)
        assert result.status == TaskStatus.FAILED
        assert agent.failure_recorded is True

        stored = await mem.retrieve("failure:t-003")
        assert stored is not None
        assert "simulated failure" in stored.value["error"]
