import unittest
from unittest.mock import Mock
from analogai.foundation.modules.agent.agent import Agent, Role
from analogai.foundation.modules.user.user import User
from analogai.foundation.infrastructure.repositories.in_memory.in_memory_agent_repository_impl import InMemoryAgentRepositoryImpl
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.mappers import AgentStorageMapper

class TestInMemoryAgentRepositoryImpl(unittest.TestCase):
    
    def setUp(self):
        self.mock_storage_adapter = Mock(spec=StorageAdapter)
        self.repository = InMemoryAgentRepositoryImpl(storage_adapter=self.mock_storage_adapter)
        
        self.test_user = User(id="user_123", name="Test User", email="test@example.com")
        self.test_agent = Agent(
            id="agent_456",
            knowledge_base="Test knowledge base",
            creator=self.test_user,
            roles=[]
        )
    
    def test_create_agent(self):
        self.repository.create(self.test_agent)
        
        expected_data = AgentStorageMapper.to_dict(self.test_agent)
        self.mock_storage_adapter.store.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            expected_data
        )
    
    def test_find_by_id_existing_agent(self):
        agent_data = AgentStorageMapper.to_dict(self.test_agent)
        self.mock_storage_adapter.retrieve_by_id.return_value = agent_data
        
        result = self.repository.find_by_id("agent_456")
        
        self.assertEqual(result.id, self.test_agent.id)
        self.assertEqual(result.knowledge_base, self.test_agent.knowledge_base)
        self.mock_storage_adapter.retrieve_by_id.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            "agent_456"
        )
    
    def test_find_by_id_nonexistent_agent(self):
        self.mock_storage_adapter.retrieve_by_id.return_value = None
        
        result = self.repository.find_by_id("nonexistent")
        
        self.assertIsNone(result)
        self.mock_storage_adapter.retrieve_by_id.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            "nonexistent"
        )
    
    def test_update_agent(self):
        self.repository.update(self.test_agent)
        
        expected_data = AgentStorageMapper.to_dict(self.test_agent)
        self.mock_storage_adapter.store.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            expected_data
        )
    
    def test_delete_agent_success(self):
        self.mock_storage_adapter.delete.return_value = True
        
        result = self.repository.delete("agent_456")
        
        self.assertTrue(result)
        self.mock_storage_adapter.delete.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            "agent_456"
        )
    
    def test_delete_agent_failure(self):
        self.mock_storage_adapter.delete.return_value = False
        
        result = self.repository.delete("nonexistent")
        
        self.assertFalse(result)
        self.mock_storage_adapter.delete.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME, 
            "nonexistent"
        )
    
    def test_find_by_creator_id_with_agents(self):
        agent_data = AgentStorageMapper.to_dict(self.test_agent)
        self.mock_storage_adapter.retrieve_by_attributes.return_value = [agent_data]
        
        result = self.repository.find_by_creator_id("user_123")
        
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].id, self.test_agent.id)
        self.mock_storage_adapter.retrieve_by_attributes.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME,
            {"creator_id": "user_123"}
        )
    
    def test_find_by_creator_id_no_agents(self):
        self.mock_storage_adapter.retrieve_by_attributes.return_value = []
        
        result = self.repository.find_by_creator_id("user_123")
        
        self.assertEqual(len(result), 0)
        self.mock_storage_adapter.retrieve_by_attributes.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME,
            {"creator_id": "user_123"}
        )
    
    def test_find_by_creator_id_multiple_agents(self):
        another_agent = Agent(
            id="agent_789",
            knowledge_base="Another knowledge base",
            creator=self.test_user,
            roles=[]
        )
        
        agent_data1 = AgentStorageMapper.to_dict(self.test_agent)
        agent_data2 = AgentStorageMapper.to_dict(another_agent)
        self.mock_storage_adapter.retrieve_by_attributes.return_value = [agent_data1, agent_data2]
        
        result = self.repository.find_by_creator_id("user_123")
        
        self.assertEqual(len(result), 2)
        ids = {r.id for r in result}
        self.assertIn(self.test_agent.id, ids)
        self.assertIn(another_agent.id, ids)
        self.mock_storage_adapter.retrieve_by_attributes.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME,
            {"creator_id": "user_123"}
        )
    
    def test_find_by_creator_id_agent_without_creator(self):
        agent_data = AgentStorageMapper.to_dict(self.test_agent)
        self.mock_storage_adapter.retrieve_by_attributes.return_value = [agent_data]
        
        result = self.repository.find_by_creator_id("user_123")
        
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].id, self.test_agent.id)
        self.mock_storage_adapter.retrieve_by_attributes.assert_called_once_with(
            InMemoryAgentRepositoryImpl.COLLECTION_NAME,
            {"creator_id": "user_123"}
        )

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