import unittest
from unittest.mock import Mock, patch
from analogai.foundation.modules.agent.agent import Agent, Role
from analogai.foundation.modules.user.user import User
from analogai.foundation.common.enums.agent_openness import AgentOpenness
from analogai.foundation.infrastructure.repositories.mongodb.mongodb_agent_repository_impl import MongoDBAgentRepositoryImpl
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.tests.repositories.mongodb.test_config import MongoDBEnvironmentCheckMixin

class TestMongoDBAgentRepositoryImpl(MongoDBEnvironmentCheckMixin, unittest.TestCase):
    
    def setUp(self):
        self.mock_storage_adapter = Mock(spec=StorageAdapter)
        self.repository = MongoDBAgentRepositoryImpl(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_find_by_id_existing_agent(self):
        doc = {
            "id": "agent_456",
            "knowledge_base": "Test knowledge base",
            "creator_id": "user_123",
            "openness": "open",
            "roles": []
        }
        self.mock_storage_adapter.retrieve_by_id.return_value = doc
        
        result = self.repository.find_by_id("agent_456")
        
        self.assertIsNotNone(result)
        self.assertEqual(result.id, "agent_456")
        self.assertEqual(result.knowledge_base, "Test knowledge base")
        self.assertEqual(result.creator.id, "user_123")
        self.mock_storage_adapter.retrieve_by_id.assert_called_once_with(
            MongoDBAgentRepositoryImpl.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(
            MongoDBAgentRepositoryImpl.COLLECTION_NAME, 
            "nonexistent"
        )
    
    def test_create_agent(self):
        self.repository.create(self.test_agent)
        
        self.mock_storage_adapter.store.assert_called_once()
        call_args = self.mock_storage_adapter.store.call_args
        self.assertEqual(call_args[0][0], MongoDBAgentRepositoryImpl.COLLECTION_NAME)
        
        stored_data = call_args[0][1]
        self.assertEqual(stored_data["id"], "agent_456")
        self.assertEqual(stored_data["knowledge_base"], "Test knowledge base")
        self.assertEqual(stored_data["creator_id"], "user_123")
    
    def test_update_agent(self):
        self.repository.update(self.test_agent)
        
        self.mock_storage_adapter.store.assert_called_once()
        call_args = self.mock_storage_adapter.store.call_args
        self.assertEqual(call_args[0][0], MongoDBAgentRepositoryImpl.COLLECTION_NAME)
    
    def test_delete_agent(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(
            MongoDBAgentRepositoryImpl.COLLECTION_NAME, 
            "agent_456"
        )
    
    def test_find_by_creator_id_with_agents(self):
        doc1 = {
            "id": "agent_456",
            "knowledge_base": "Test knowledge base",
            "creator_id": "user_123",
            "openness": "open",
            "roles": []
        }
        
        self.mock_storage_adapter.retrieve_by_attributes.return_value = [doc1]
        
        result = self.repository.find_by_creator_id("user_123")
        
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].id, "agent_456")
        self.assertEqual(result[0].creator.id, "user_123")
        self.mock_storage_adapter.retrieve_by_attributes.assert_called_once_with(
            MongoDBAgentRepositoryImpl.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(
            MongoDBAgentRepositoryImpl.COLLECTION_NAME,
            {"creator_id": "user_123"}
        )
    

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