import unittest
from pygeai_orchestration import (
    Message,
    MessageRole,
    Conversation,
    Context,
    ContextManager,
    State,
    StateStatus,
    StateManager,
    Memory,
    MemoryStore
)


class TestMessage(unittest.TestCase):

    def test_creation(self):
        msg = Message(role=MessageRole.USER, content="Hello")

        self.assertEqual(msg.role, MessageRole.USER)
        self.assertEqual(msg.content, "Hello")
        self.assertIsNotNone(msg.timestamp)

    def test_with_metadata(self):
        msg = Message(
            role=MessageRole.ASSISTANT,
            content="Response",
            metadata={"key": "value"}
        )

        self.assertEqual(msg.metadata["key"], "value")


class TestConversation(unittest.TestCase):

    def test_creation(self):
        conv = Conversation(id="conv-1")

        self.assertEqual(conv.id, "conv-1")
        self.assertEqual(len(conv.messages), 0)

    def test_add_message(self):
        conv = Conversation(id="conv-1")
        msg = Message(role=MessageRole.USER, content="Hello")

        conv.add_message(msg)
        self.assertEqual(len(conv.messages), 1)

    def test_get_messages_by_role(self):
        conv = Conversation(id="conv-1")
        conv.add_message(Message(role=MessageRole.USER, content="Q1"))
        conv.add_message(Message(role=MessageRole.ASSISTANT, content="A1"))
        conv.add_message(Message(role=MessageRole.USER, content="Q2"))

        user_msgs = conv.get_messages(MessageRole.USER)
        self.assertEqual(len(user_msgs), 2)

    def test_to_dict_list(self):
        conv = Conversation(id="conv-1")
        conv.add_message(Message(role=MessageRole.USER, content="Hello"))

        dict_list = conv.to_dict_list()
        self.assertEqual(len(dict_list), 1)
        self.assertEqual(dict_list[0]["role"], "user")


class TestContext(unittest.TestCase):

    def test_creation(self):
        ctx = Context()
        self.assertEqual(len(ctx.data), 0)

    def test_set_get(self):
        ctx = Context()
        ctx.set("key", "value")

        self.assertTrue(ctx.has("key"))
        self.assertEqual(ctx.get("key"), "value")

    def test_remove(self):
        ctx = Context()
        ctx.set("key", "value")
        ctx.remove("key")

        self.assertFalse(ctx.has("key"))

    def test_update(self):
        ctx = Context()
        ctx.update({"a": 1, "b": 2})

        self.assertEqual(ctx.get("a"), 1)
        self.assertEqual(ctx.get("b"), 2)


class TestContextManager(unittest.TestCase):

    def test_create_context(self):
        manager = ContextManager()
        ctx = manager.create_context("ctx-1")

        self.assertIsNotNone(ctx)
        self.assertIn("ctx-1", manager.list_contexts())

    def test_get_context(self):
        manager = ContextManager()
        manager.create_context("ctx-1", {"test": "data"})

        ctx = manager.get_context("ctx-1")
        self.assertEqual(ctx.get("test"), "data")

    def test_delete_context(self):
        manager = ContextManager()
        manager.create_context("ctx-1")
        manager.delete_context("ctx-1")

        self.assertNotIn("ctx-1", manager.list_contexts())


class TestState(unittest.TestCase):

    def test_creation(self):
        state = State()
        self.assertEqual(state.status, StateStatus.INITIALIZED)

    def test_update_status(self):
        state = State()
        state.update_status(StateStatus.RUNNING)

        self.assertEqual(state.status, StateStatus.RUNNING)

    def test_checkpoint(self):
        state = State()
        state.set("key", "value")
        state.create_checkpoint()

        state.set("key", "new_value")
        self.assertEqual(state.get("key"), "new_value")

        state.restore_checkpoint()
        self.assertEqual(state.get("key"), "value")


class TestStateManager(unittest.TestCase):

    def test_create_state(self):
        manager = StateManager()
        state = manager.create_state("state-1")

        self.assertIsNotNone(state)
        self.assertIn("state-1", manager.list_states())

    def test_update_state(self):
        manager = StateManager()
        manager.create_state("state-1")

        success = manager.update_state("state-1", StateStatus.COMPLETED, {"result": "done"})
        self.assertTrue(success)

        state = manager.get_state("state-1")
        self.assertEqual(state.status, StateStatus.COMPLETED)


class TestMemory(unittest.TestCase):

    def test_creation(self):
        mem = Memory()
        self.assertEqual(mem.size(), 0)

    def test_store_retrieve(self):
        mem = Memory()
        mem.store("fact1", "The sky is blue")

        self.assertTrue(mem.has("fact1"))
        self.assertEqual(mem.retrieve("fact1"), "The sky is blue")

    def test_max_size(self):
        mem = Memory(max_size=2)
        mem.store("k1", "v1")
        mem.store("k2", "v2")
        mem.store("k3", "v3")

        self.assertEqual(mem.size(), 2)
        self.assertFalse(mem.has("k1"))

    def test_get_recent(self):
        mem = Memory()
        mem.store("k1", "v1")
        mem.store("k2", "v2")
        mem.store("k3", "v3")

        recent = mem.get_recent(limit=2)
        self.assertEqual(len(recent), 2)


class TestMemoryStore(unittest.TestCase):

    def test_create_memory(self):
        store = MemoryStore()
        mem = store.create_memory("mem-1")

        self.assertIsNotNone(mem)
        self.assertIn("mem-1", store.list_memories())

    def test_get_memory(self):
        store = MemoryStore()
        store.create_memory("mem-1")

        mem = store.get_memory("mem-1")
        self.assertIsNotNone(mem)


if __name__ == '__main__':
    unittest.main()
