import unittest
from unittest import TestCase
from unittest.mock import Mock
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement as Statement
from analogai.foundation.core.value_objects.relation import Relation

from analogai.foundation.modules.user.user import User
from analogai.foundation.modules.agent.agent import Agent
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
from analogai.foundation.settings import config

class TestInMemoryBeliefRepositoryImpl(TestCase):
    def setUp(self):
        self._config_override = config.override(**{"vector.enabled": False})
        self._config_override.__enter__()
        self.storage_adapter = Mock(spec=StorageAdapter)
        self.repository = InMemoryBeliefRepositoryImpl(storage_adapter=self.storage_adapter)
        self.user = User(id="user_id", name="Test User", email="test@example.com")
        self.agent = Agent(id="agent_id", knowledge_base="Test KB", creator=self.user, roles=[])
        self.subject_notion = Notion(user_id=self.user.id, agent_id=self.agent.id, caption="human", id="human_id")
        self.complement_notion = Notion(user_id=self.user.id, agent_id=self.agent.id, caption="mortal", id="mortal_id")
        self.statement = 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,
            id="stmt_id"
        )
        self.belief = Belief(agent_id=self.agent.id, id="belief_id")
        self.belief.add_statement(self.statement)

    def test_save_calls_store(self):
        saved = self.repository.save(self.belief)
        
        self.storage_adapter.store.assert_called_once()
        args, _ = self.storage_adapter.store.call_args
        self.assertEqual(args[0], "beliefs")
        self.assertEqual(saved.subject_notion.id, self.subject_notion.id)
        self.assertEqual(saved.complement_notion.id, self.complement_notion.id)

    def test_find(self):
        belief_data = BeliefStorageMapper.to_dict(self.belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

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

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect

        result = self.repository.find(self.subject_notion, self.complement_notion, Relation.from_string("is"))
        
        expected_query = {
            'subject_notion.id': "human_id",
            'complement_notion.id': "mortal_id",
            'relation': "is"
        }
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(result.id, self.belief.id)

    def test_find_no_result(self):
        self.storage_adapter.retrieve_by_attributes.return_value = []
        
        result = self.repository.find(self.subject_notion, self.complement_notion, Relation.from_string("is"))
        
        expected_query = {
            'subject_notion.id': "human_id",
            'complement_notion.id': "mortal_id",
            'relation': "is"
        }
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertIsNone(result)

    def test_find_by_id(self):
        belief_data = BeliefStorageMapper.to_dict(self.belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        
        def retrieve_by_id_side_effect(collection, id):
            if collection == "beliefs":
                return belief_data
            elif collection == "statements":
                return statement_data
            return None
        
        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect
        
        result = self.repository.find_by_id("belief_id")
        
        self.assertEqual(self.storage_adapter.retrieve_by_id.call_count, 2)
        self.storage_adapter.retrieve_by_id.assert_any_call("beliefs", "belief_id")
        self.storage_adapter.retrieve_by_id.assert_any_call("statements", "stmt_id")
        self.assertEqual(result.id, self.belief.id)

    def test_find_by_id_not_found(self):
        self.storage_adapter.retrieve_by_id.return_value = None
        
        result = self.repository.find_by_id("nonexistent_id")
        
        self.storage_adapter.retrieve_by_id.assert_called_once_with("beliefs", "nonexistent_id")
        self.assertIsNone(result)

    def tearDown(self):
        if self._config_override:
            self._config_override.__exit__(None, None, None)

    def test_delete(self):
        self.repository.delete(self.belief)
        
        self.storage_adapter.delete.assert_called_once_with("beliefs", "belief_id")

    def test_search_by_notions(self):
        belief_data = BeliefStorageMapper.to_dict(self.belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

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

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect

        results = self.repository.search_by_notions(self.subject_notion, self.complement_notion)
        
        expected_query = {
            'subject_notion.id': "human_id",
            'complement_notion.id': "mortal_id"
        }
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].id, self.belief.id)

    def test_search_by_subject_notion(self):
        belief_data = BeliefStorageMapper.to_dict(self.belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

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

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect

        results = self.repository.search_by_subject_notion(self.subject_notion)
        
        expected_query = {'subject_notion.id': "human_id"}
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].id, self.belief.id)

    def test_search_by_complement_notion(self):
        belief_data = BeliefStorageMapper.to_dict(self.belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        self.storage_adapter.retrieve_by_attributes.return_value = [belief_data]

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

        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect

        results = self.repository.search_by_complement_notion(self.complement_notion)
        
        expected_query = {'complement_notion.id': "mortal_id"}
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].id, self.belief.id)

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