import unittest
from unittest.mock import Mock
from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement as Statement

from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.time import Time
from analogai.foundation.common.enums.deontic_modality import DeonticModality
from analogai.foundation.infrastructure.repositories.in_memory.in_memory_belief_repository_impl import InMemoryBeliefRepositoryImpl
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.mappers import BeliefStorageMapper, StatementStorageMapper

class TestBeliefRepositoryModifierFiltering(unittest.TestCase):
    
    def setUp(self):
        self.storage_adapter = Mock(spec=StorageAdapter)
        self.repository = InMemoryBeliefRepositoryImpl(storage_adapter=self.storage_adapter)
        self.user_id = "user1"
        self.agent_id = "agent1"
        self.subject_notion = Notion(user_id=self.user_id, agent_id=self.agent_id, caption="Source", id="source1")
        self.complement_notion = Notion(user_id=self.user_id, agent_id=self.agent_id, caption="Target", id="target1")
    
    def create_statement(self, location=None, from_time=None, to_time=None, deontic_modality=None):
        modifier = None
        if any([location, from_time, to_time, deontic_modality]):
            from analogai.foundation.common.utils.modifier_utils import ModifierBuilder
            builder = ModifierBuilder()
            if location:
                builder.with_location(location)
            if from_time:
                builder.with_from_time(
                    year=from_time.year,
                    month=from_time.month,
                    day=from_time.day,
                    hour=from_time.hour,
                    minute=from_time.minute
                )
            if to_time:
                builder.with_to_time(
                    year=to_time.year,
                    month=to_time.month,
                    day=to_time.day,
                    hour=to_time.hour,
                    minute=to_time.minute
                )
            if deontic_modality:
                builder.with_deontic_modality(deontic_modality)
            modifier = builder.build()
        return Statement(
            user_id=self.user_id,
            agent_id=self.agent_id,
            subject_notion=self.subject_notion,
            relation=Relation.from_string("is"),
            complement_notion=self.complement_notion,
            probability=0.8,
            validity=0.9,
            modifier=modifier
        )
    
    def test_find_with_matching_statement_modifier(self):
        statement = self.create_statement(location="New York")
        belief = Belief(agent_id=self.agent_id)
        belief.add_statement(statement)
        belief_data = BeliefStorageMapper.to_dict(belief)
        statement_data = StatementStorageMapper.to_dict(statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

        def retrieve_by_id_side_effect(collection, id):
            if collection == "statements" and id == statement.id:
                return statement_data
            return None

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect
        modifier = Modifier(location="New York")
        
        result = self.repository.find(
            self.subject_notion,
            self.complement_notion,
            Relation.from_string("is"),
            modifier=modifier
        )
        
        self.assertIsNotNone(result)
        self.assertEqual(len(result.statements), 1)
        self.assertEqual(result.statements[0].location, "New York")
    
    def test_find_with_non_matching_statement_modifier(self):
        statement = self.create_statement(location="New York")
        belief = Belief(agent_id=self.agent_id)
        belief.add_statement(statement)
        belief_data = BeliefStorageMapper.to_dict(belief)
        statement_data = StatementStorageMapper.to_dict(statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

        def retrieve_by_id_side_effect(collection, id):
            if collection == "statements" and id == statement.id:
                return statement_data
            return None

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect
        modifier = Modifier(location="London")
        
        result = self.repository.find(
            self.subject_notion,
            self.complement_notion,
            Relation.from_string("is"),
            modifier=modifier
        )
        
        self.assertIsNone(result)
    
    def test_search_by_notions_with_modifier(self):
        statement1 = self.create_statement(location="New York", deontic_modality=DeonticModality.OBLIGATORY)
        statement2 = self.create_statement(location="London")
        belief1 = Belief(agent_id=self.agent_id)
        belief1.add_statement(statement1)
        belief2 = Belief(agent_id=self.agent_id)
        belief2.add_statement(statement2)
        belief1_data = BeliefStorageMapper.to_dict(belief1)
        belief2_data = BeliefStorageMapper.to_dict(belief2)
        statement1_data = StatementStorageMapper.to_dict(statement1)
        statement2_data = StatementStorageMapper.to_dict(statement2)
        
        def retrieve_by_id_side_effect(collection, id):
            if collection == "statements":
                if id == statement1.id:
                    return statement1_data
                elif id == statement2.id:
                    return statement2_data
            return None
        
        self.storage_adapter.retrieve_by_attributes.return_value = [belief1_data, belief2_data]
        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect
        modifier = Modifier(location="New York")
        
        results = self.repository.search_by_notions(
            self.subject_notion,
            self.complement_notion,
            modifier=modifier
        )
        
        self.assertEqual(len(results), 1)
        self.assertEqual(len(results[0].statements), 1)
        self.assertEqual(results[0].statements[0].location, "New York")
    
    def test_search_by_notions_with_modifier(self):
        statement1 = self.create_statement(location="New York")
        statement2 = self.create_statement(location="London")
        belief = Belief(agent_id=self.agent_id)
        belief.add_statement(statement1)
        belief.add_statement(statement2)
        belief_data = BeliefStorageMapper.to_dict(belief)
        statement1_data = StatementStorageMapper.to_dict(statement1)
        statement2_data = StatementStorageMapper.to_dict(statement2)
        
        def retrieve_by_id_side_effect(collection, id):
            if collection == "statements":
                if id == statement1.id:
                    return statement1_data
                elif id == statement2.id:
                    return statement2_data
            return None
        
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]
        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect
        modifier = Modifier(location="New York")
        
        results = self.repository.search_by_notions(
            self.subject_notion,
            self.complement_notion,
            modifier=modifier
        )
        
        self.assertEqual(len(results), 1)
        self.assertEqual(len(results[0].statements), 2)

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

