import types
from unittest import IsolatedAsyncioTestCase, main
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4

from avalan.entities import EngineMessage, Message, MessageRole
from avalan.event import EventType
from avalan.event.manager import EventManager
from avalan.memory import RecentMessageMemory
from avalan.memory.manager import MemoryManager
from avalan.memory.permanent import (
    PermanentMemory,
    PermanentMemoryStore,
    PermanentMessageMemory,
    VectorFunction,
)


class MemoryManagerCreateTestCase(IsolatedAsyncioTestCase):
    async def test_create_instance_with_recent_only(self):
        tp = AsyncMock()
        agent_id = uuid4()
        participant_id = uuid4()

        logger = MagicMock()
        manager = await MemoryManager.create_instance(
            agent_id=agent_id,
            participant_id=participant_id,
            text_partitioner=tp,
            logger=logger,
            event_manager=MagicMock(spec=EventManager),
        )

        self.assertIsInstance(manager.recent_message, RecentMessageMemory)
        self.assertIsNone(manager.permanent_message)
        self.assertTrue(manager.has_recent_message)
        self.assertFalse(manager.has_permanent_message)
        self.assertIs(manager._logger, logger)

    async def test_create_instance_with_permanent(self):
        tp = AsyncMock()
        pmemory = MagicMock(spec=PermanentMessageMemory)
        agent_id = uuid4()
        participant_id = uuid4()

        with self.subTest():
            dummy = types.SimpleNamespace()

            class PgsqlDummy:
                pass

            PgsqlDummy.create_instance = AsyncMock(return_value=pmemory)
            dummy.PgsqlMessageMemory = PgsqlDummy
            with patch.dict(
                "sys.modules", {"avalan.memory.permanent.pgsql.message": dummy}
            ):
                logger = MagicMock()
                manager = await MemoryManager.create_instance(
                    agent_id=agent_id,
                    participant_id=participant_id,
                    text_partitioner=tp,
                    logger=logger,
                    with_permanent_message_memory="dsn",
                    event_manager=MagicMock(spec=EventManager),
                )
            PgsqlDummy.create_instance.assert_awaited_once_with(
                dsn="dsn", logger=logger
            )

        self.assertIs(manager.permanent_message, pmemory)
        self.assertIsInstance(manager.recent_message, RecentMessageMemory)
        self.assertTrue(manager.has_recent_message)
        self.assertTrue(manager.has_permanent_message)
        self.assertIs(manager._logger, logger)


class MemoryManagerOperationTestCase(IsolatedAsyncioTestCase):
    def setUp(self):
        self.tp = AsyncMock()
        self.pm = AsyncMock(spec=PermanentMessageMemory)
        self.rm = RecentMessageMemory()
        self.event_manager = MagicMock()
        self.event_manager.trigger = AsyncMock()
        self.manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=self.pm,
            recent_message_memory=self.rm,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=self.event_manager,
        )
        self.permanent = AsyncMock(spec=PermanentMemory)

    async def test_append_message(self):
        partitions = ["p"]
        self.tp.return_value = partitions
        msg = EngineMessage(
            agent_id=self.manager._agent_id,
            model_id="m",
            message=Message(role=MessageRole.USER, content="hi"),
        )

        await self.manager.append_message(msg)

        self.tp.assert_awaited_once_with("hi")
        self.pm.append_with_partitions.assert_awaited_once_with(
            msg, partitions=partitions
        )
        self.assertEqual(self.rm.data, [msg])

    async def test_continue_session_load_recent(self):
        messages = [
            EngineMessage(
                agent_id=self.manager._agent_id,
                model_id="m",
                message=Message(role=MessageRole.USER, content="x"),
            )
        ]
        self.pm.get_recent_messages.return_value = messages

        await self.manager.continue_session(uuid4())

        self.pm.continue_session.assert_awaited()
        self.pm.get_recent_messages.assert_awaited()
        self.assertEqual(self.rm.data, messages)

    async def test_start_session(self):
        await self.manager.start_session()
        self.pm.reset_session.assert_awaited_once()
        self.assertTrue(self.rm.is_empty)

    async def test_search_messages(self):
        partitions = ["p1"]
        self.tp.return_value = partitions
        result = [
            EngineMessage(
                agent_id=self.manager._agent_id,
                model_id="m",
                message=Message(role=MessageRole.USER, content="hi"),
            )
        ]
        self.pm.search_messages.return_value = result

        messages = await self.manager.search_messages(
            "hi",
            agent_id=uuid4(),
            session_id=uuid4(),
            participant_id=uuid4(),
            function=VectorFunction.L2_DISTANCE,
        )

        self.tp.assert_awaited_once_with("hi")
        self.pm.search_messages.assert_awaited_once()
        self.assertEqual(messages, result)

    async def test_search_permanent_memory(self):
        namespace = "docs"
        partitions = ["p2"]
        participant_id = uuid4()
        self.tp.return_value = partitions
        self.permanent.search_memories.return_value = ["memory"]
        self.manager.add_permanent_memory(namespace, self.permanent)

        memories = await self.manager.search_partitions(
            "query",
            participant_id=participant_id,
            namespace=namespace,
            function=VectorFunction.L2_DISTANCE,
            limit=5,
        )

        self.tp.assert_awaited_once_with("query")
        self.permanent.search_memories.assert_awaited_once_with(
            search_partitions=partitions,
            participant_id=participant_id,
            namespace=namespace,
            function=VectorFunction.L2_DISTANCE,
            limit=5,
        )
        self.assertEqual(memories, ["memory"])

    async def test_search_permanent_memory_missing_namespace(self):
        with self.assertRaises(KeyError) as caught:
            await self.manager.search_partitions(
                "query",
                participant_id=uuid4(),
                namespace="missing",
                function=VectorFunction.L2_DISTANCE,
            )

        self.assertIn(
            "Memory namespace missing not defined", str(caught.exception)
        )
        self.tp.assert_not_called()

    async def test_late_message_memory_uses_lazy_partitioner(self):
        text_partitioner = AsyncMock(return_value=["p"])
        factory = MagicMock(return_value=text_partitioner)
        permanent_message_memory = AsyncMock(spec=PermanentMessageMemory)
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=None,
            text_partitioner_factory=factory,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )

        manager.add_permanent_message_memory(permanent_message_memory)
        factory.assert_not_called()
        msg = EngineMessage(
            agent_id=manager._agent_id,
            model_id="m",
            message=Message(role=MessageRole.USER, content="hi"),
        )
        await manager.append_message(msg)

        factory.assert_called_once_with()
        text_partitioner.assert_awaited_once_with("hi")
        permanent_message_memory.append_with_partitions.assert_awaited_once_with(
            msg, partitions=["p"]
        )

    async def test_late_permanent_memory_uses_lazy_partitioner_on_search(self):
        namespace = "docs"
        text_partitioner = AsyncMock(return_value=["p"])
        factory = MagicMock(return_value=text_partitioner)
        permanent_memory = AsyncMock(spec=PermanentMemory)
        permanent_memory.search_memories.return_value = ["memory"]
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=None,
            text_partitioner_factory=factory,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )

        manager.add_permanent_memory(namespace, permanent_memory)
        factory.assert_not_called()
        memories = await manager.search_partitions(
            "query",
            participant_id=uuid4(),
            namespace=namespace,
            function=VectorFunction.L2_DISTANCE,
        )

        factory.assert_called_once_with()
        text_partitioner.assert_awaited_once_with("query")
        self.assertEqual(memories, ["memory"])

    async def test_search_permanent_memory_requires_partitioner(self):
        namespace = "docs"
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=None,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        manager.add_permanent_memory(namespace, self.permanent)

        with self.assertRaisesRegex(
            AssertionError,
            "Text partitioner is required to search permanent memory",
        ):
            await manager.search_partitions(
                "query",
                participant_id=uuid4(),
                namespace=namespace,
                function=VectorFunction.L2_DISTANCE,
            )

    async def test_list_memories_missing_namespace(self):
        with self.assertRaises(KeyError) as caught:
            await self.manager.list_memories(
                participant_id=uuid4(),
                namespace="missing",
            )

        self.assertIn(
            "Memory namespace missing not defined", str(caught.exception)
        )
        self.permanent.list_memories.assert_not_called()

    def test_add_and_delete_permanent_memory(self):
        self.manager.add_permanent_memory("code", self.permanent)
        self.assertIn("code", self.manager._permanent_memory_stores)
        store = self.manager._permanent_memory_stores["code"]
        self.assertIs(store[0], self.permanent)
        self.assertIsNone(store[1])
        self.manager.delete_permanent_memory("code")
        self.assertNotIn("code", self.manager._permanent_memory_stores)

    async def test_list_permanent_memories(self):
        namespace = "docs"
        participant_id = uuid4()
        expected = [MagicMock()]
        self.permanent.list_memories.return_value = expected
        self.manager.add_permanent_memory(namespace, self.permanent)

        memories = await self.manager.list_memories(
            participant_id=participant_id,
            namespace=namespace,
        )

        self.permanent.list_memories.assert_awaited_once_with(
            participant_id=participant_id,
            namespace=namespace,
        )
        self.assertIs(memories, expected)

    def test_list_permanent_memory_stores(self):
        stores = self.manager.list_permanent_memory_stores()
        self.assertEqual(stores, [])
        self.manager.add_permanent_memory(
            "code", self.permanent, description="Code docs"
        )
        stores = self.manager.list_permanent_memory_stores()
        self.assertEqual(len(stores), 1)
        self.assertIsInstance(stores[0], PermanentMemoryStore)
        self.assertEqual(stores[0].namespace, "code")
        self.assertEqual(stores[0].description, "Code docs")


class MemoryManagerInitTestCase(IsolatedAsyncioTestCase):
    def test_constructor_variants(self):
        agent_id = uuid4()
        participant_id = uuid4()
        tp = AsyncMock()

        manager = MemoryManager(
            agent_id=agent_id,
            participant_id=participant_id,
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertFalse(manager.has_recent_message)
        self.assertFalse(manager.has_permanent_message)
        self.assertEqual(manager.participant_id, participant_id)
        self.assertIsNone(manager.recent_messages)

        pmem = MagicMock(spec=PermanentMemory)
        manager = MemoryManager(
            agent_id=agent_id,
            participant_id=participant_id,
            permanent_message_memory=None,
            recent_message_memory=RecentMessageMemory(),
            text_partitioner=tp,
            logger=MagicMock(),
            permanent_memory_stores={
                "code": (pmem, "Docs"),
            },
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertTrue(manager.has_recent_message)
        self.assertIn("code", manager._permanent_memory_stores)
        self.assertEqual(manager._permanent_memory_stores["code"][1], "Docs")

        pmem2 = MagicMock(spec=PermanentMemory)
        manager = MemoryManager(
            agent_id=agent_id,
            participant_id=participant_id,
            permanent_message_memory=None,
            recent_message_memory=RecentMessageMemory(),
            text_partitioner=tp,
            logger=MagicMock(),
            permanent_memory_stores={
                "code": (pmem, "Docs"),
                "docs": (pmem2, "Documents"),
            },
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertEqual(len(manager._permanent_memory_stores), 2)
        self.assertEqual(
            manager._permanent_memory_stores["docs"][1], "Documents"
        )


class MemoryManagerPropertyTestCase(IsolatedAsyncioTestCase):
    async def test_recent_messages_property(self):
        tp = AsyncMock()
        rm = RecentMessageMemory()
        msg = EngineMessage(
            agent_id=uuid4(),
            model_id="m",
            message=Message(role=MessageRole.USER, content="hi"),
        )
        await rm.append(msg.agent_id, msg)
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=rm,
            text_partitioner=tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertEqual(manager.recent_messages, [msg])
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertIsNone(manager.recent_messages)


class MemoryManagerMethodsTestCase(IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.tp = AsyncMock()
        self.pm = AsyncMock(spec=PermanentMessageMemory)
        self.rm = RecentMessageMemory()

    async def test_append_message_variants(self):
        msg = EngineMessage(
            agent_id=uuid4(),
            model_id="m",
            message=Message(role=MessageRole.USER, content="hi"),
        )

        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=self.pm,
            recent_message_memory=None,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.append_message(msg)
        self.tp.assert_awaited_once()
        self.pm.append_with_partitions.assert_awaited_once()

        self.tp.reset_mock()
        self.pm.reset_mock()
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=self.rm,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.append_message(msg)
        self.tp.assert_not_called()
        self.assertEqual(manager.recent_messages, [msg])

        self.tp.reset_mock()
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.append_message(msg)
        self.tp.assert_not_called()

    async def test_continue_and_start_session_variants(self):
        session_id = uuid4()
        messages = [
            EngineMessage(
                agent_id=uuid4(),
                model_id="m",
                message=Message(role=MessageRole.USER, content="x"),
            )
        ]
        self.pm.get_recent_messages.return_value = messages

        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=self.pm,
            recent_message_memory=None,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.continue_session(session_id)
        self.pm.continue_session.assert_awaited()
        self.pm.get_recent_messages.assert_not_awaited()
        await manager.start_session()
        self.pm.reset_session.assert_awaited()

        self.pm.reset_mock()
        self.pm.get_recent_messages.reset_mock()
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=self.rm,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.continue_session(session_id)
        self.pm.continue_session.assert_not_awaited()
        self.assertEqual(manager.recent_messages, [])
        await manager.start_session()
        self.assertTrue(self.rm.is_empty)

        self.pm.reset_mock()
        self.pm.get_recent_messages.reset_mock()
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        await manager.continue_session(session_id)
        await manager.start_session()


class MemoryManagerContextTestCase(IsolatedAsyncioTestCase):
    async def test_context_exit(self):
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=None,
            text_partitioner=AsyncMock(),
            logger=MagicMock(),
            event_manager=MagicMock(spec=EventManager),
        )
        self.assertIsNone(manager.__exit__(None, None, None))


class MemoryManagerEventTestCase(IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.tp = AsyncMock()
        self.pm = AsyncMock(spec=PermanentMessageMemory)
        self.pm.session_id = uuid4()
        self.rm = RecentMessageMemory()
        self.event_manager = MagicMock(spec=EventManager)
        self.event_manager.trigger = AsyncMock()
        self.manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=self.pm,
            recent_message_memory=self.rm,
            text_partitioner=self.tp,
            logger=MagicMock(),
            event_manager=self.event_manager,
        )

    async def test_append_message_triggers_events(self):
        msg = EngineMessage(
            agent_id=self.manager._agent_id,
            model_id="m",
            message=Message(role=MessageRole.USER, content="hi"),
        )

        await self.manager.append_message(msg)

        called_types = [
            c.args[0].type for c in self.event_manager.trigger.await_args_list
        ]
        self.assertIn(EventType.MEMORY_PERMANENT_MESSAGE_ADD, called_types)
        self.assertIn(EventType.MEMORY_PERMANENT_MESSAGE_ADDED, called_types)
        self.assertTrue(
            any(
                c.args[0].type == EventType.MEMORY_PERMANENT_MESSAGE_ADDED
                and c.args[0].elapsed is not None
                for c in self.event_manager.trigger.await_args_list
            )
        )
        self.assertTrue(
            any(
                c.args[0].type == EventType.MEMORY_PERMANENT_MESSAGE_ADD
                and c.args[0].payload.get("participant_id")
                == self.manager.participant_id
                for c in self.event_manager.trigger.await_args_list
            )
        )

    async def test_continue_session_triggers_events(self):
        session_id = uuid4()
        await self.manager.continue_session(session_id)

        called_types = [
            c.args[0].type for c in self.event_manager.trigger.await_args_list
        ]
        self.assertIn(
            EventType.MEMORY_PERMANENT_MESSAGE_SESSION_CONTINUE, called_types
        )
        self.assertIn(
            EventType.MEMORY_PERMANENT_MESSAGE_SESSION_CONTINUED, called_types
        )
        self.assertTrue(
            any(
                c.args[0].type
                == EventType.MEMORY_PERMANENT_MESSAGE_SESSION_CONTINUED
                and c.args[0].elapsed is not None
                for c in self.event_manager.trigger.await_args_list
            )
        )
        self.assertTrue(
            any(
                c.args[0].type
                == EventType.MEMORY_PERMANENT_MESSAGE_SESSION_CONTINUE
                and c.args[0].payload.get("participant_id")
                == self.manager.participant_id
                for c in self.event_manager.trigger.await_args_list
            )
        )

    async def test_start_session_triggers_events(self):
        await self.manager.start_session()

        called_types = [
            c.args[0].type for c in self.event_manager.trigger.await_args_list
        ]
        self.assertIn(
            EventType.MEMORY_PERMANENT_MESSAGE_SESSION_START, called_types
        )
        self.assertIn(
            EventType.MEMORY_PERMANENT_MESSAGE_SESSION_STARTED, called_types
        )
        self.assertTrue(
            any(
                c.args[0].type
                == EventType.MEMORY_PERMANENT_MESSAGE_SESSION_STARTED
                and c.args[0].elapsed is not None
                for c in self.event_manager.trigger.await_args_list
            )
        )
        self.assertTrue(
            any(
                c.args[0].type
                == EventType.MEMORY_PERMANENT_MESSAGE_SESSION_START
                and c.args[0].payload.get("participant_id")
                == self.manager.participant_id
                for c in self.event_manager.trigger.await_args_list
            )
        )


class MemoryManagerEdgeCaseTestCase(IsolatedAsyncioTestCase):
    async def test_append_message_skips_non_engine_message(self):
        manager = MemoryManager(
            agent_id=uuid4(),
            participant_id=uuid4(),
            permanent_message_memory=None,
            recent_message_memory=RecentMessageMemory(),
            text_partitioner=AsyncMock(),
            logger=MagicMock(),
            event_manager=None,
        )

        await manager.append_message("bad")

        manager._logger.info.assert_called_once_with(
            "Skipping non engine message %s", "bad"
        )
        manager._text_partitioner.assert_not_called()
        self.assertTrue(manager.recent_message.is_empty)


if __name__ == "__main__":
    main()
