import unittest
from unittest import TestCase
from unittest.mock import Mock
from datetime import datetime, timedelta
from time import sleep
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 copy import copy
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_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, doc_id):
            if collection == "statements" and doc_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, doc_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 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, doc_id):
            if collection == "statements" and doc_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, doc_id):
            if collection == "statements" and doc_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, doc_id):
            if collection == "statements" and doc_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)

    def test_find_by_agent_id(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, doc_id):
            if collection == "statements" and doc_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.find_by_agent_id("agent_id")
        
        expected_query = {'agent_id': "agent_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_find_by_agent_id_no_results(self):
        self.storage_adapter.retrieve_by_attributes.return_value = []
        
        results = self.repository.find_by_agent_id("nonexistent_agent_id")
        
        expected_query = {'agent_id': "nonexistent_agent_id"}
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(results, [])

    def test_belief_has_created_at_timestamp(self):
        self.assertTrue(hasattr(self.belief, 'created_at'))
        self.assertIsInstance(self.belief.created_at, datetime)

    def test_belief_has_updated_at_timestamp(self):
        self.assertTrue(hasattr(self.belief, 'updated_at'))
        self.assertIsInstance(self.belief.updated_at, datetime)

    def test_belief_timestamps_are_set_on_creation(self):
        self.assertIsNotNone(self.belief.created_at)
        self.assertIsNotNone(self.belief.updated_at)
        self.assertEqual(self.belief.created_at, self.belief.updated_at)

    def test_belief_updated_at_changes_when_statement_added(self):
        created_at = self.belief.created_at
        updated_at = self.belief.updated_at
        
        new_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.7,
            validity=0.8,
            id="stmt_id_2"
        )
        
        sleep(0.001)
        self.belief.add_statement(new_statement)
        
        self.assertEqual(self.belief.created_at, created_at)
        self.assertGreater(self.belief.updated_at, updated_at)

    def test_repository_stores_belief_with_timestamps(self):
        belief = Belief(agent_id=self.agent.id)
        belief.add_statement(self.statement)
        
        self.repository.save(belief)
        
        self.assertEqual(self.storage_adapter.store.call_count, 1)
        stored_data = self.storage_adapter.store.call_args[0][1]
        self.assertIn('created_at', stored_data)
        self.assertIn('updated_at', stored_data)
        self.assertIsNotNone(stored_data.get('created_at'))
        self.assertIsNotNone(stored_data.get('updated_at'))

    def test_get_recently_updated_returns_recent_beliefs(self):
        recent_belief = Belief(agent_id=self.agent.id, id="recent_belief_id")
        recent_belief.add_statement(self.statement)
        recent_belief.created_at = datetime.now()
        recent_belief.updated_at = datetime.now()
        
        old_statement = Statement(
            user_id=self.user.id,
            agent_id=self.agent.id,
            subject_notion=self.subject_notion,
            relation=Relation.from_string("causes"),
            complement_notion=self.complement_notion,
            probability=0.8,
            validity=0.9,
            id="old_stmt_id"
        )
        old_belief = Belief(agent_id=self.agent.id, id="old_belief_id")
        old_belief.add_statement(old_statement)
        old_belief.created_at = datetime.now() - timedelta(minutes=2)
        old_belief.updated_at = datetime.now() - timedelta(minutes=2)
        
        recent_belief_data = BeliefStorageMapper.to_dict(recent_belief)
        old_belief_data = BeliefStorageMapper.to_dict(old_belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)
        old_statement_data = StatementStorageMapper.to_dict(old_statement)

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

        self.storage_adapter.retrieve_by_attributes.return_value = [recent_belief_data, old_belief_data]
        self.storage_adapter.retrieve_by_id.side_effect = retrieve_by_id_side_effect

        one_minute_ago = datetime.now() - timedelta(minutes=1)
        results = self.repository.get_recently_updated(self.agent.id, one_minute_ago)
        
        expected_query = {'agent_id': "agent_id"}
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].id, "recent_belief_id")

    def test_get_recently_updated_returns_beliefs_updated_recently(self):
        belief = Belief(agent_id=self.agent.id, id="updated_belief_id")
        belief.add_statement(self.statement)
        belief.created_at = datetime.now() - timedelta(minutes=2)
        belief.updated_at = datetime.now()
        
        belief_data = BeliefStorageMapper.to_dict(belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)

        def retrieve_by_id_side_effect(collection, doc_id):
            if collection == "statements" and doc_id == self.statement.id:
                return statement_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

        one_minute_ago = datetime.now() - timedelta(minutes=1)
        results = self.repository.get_recently_updated(self.agent.id, one_minute_ago)
        
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0].id, "updated_belief_id")

    def test_get_recently_updated_excludes_old_beliefs(self):
        old_belief = Belief(agent_id=self.agent.id, id="old_belief_id")
        old_belief.add_statement(self.statement)
        old_belief.created_at = datetime.now() - timedelta(minutes=2)
        old_belief.updated_at = datetime.now() - timedelta(minutes=2)
        
        belief_data = BeliefStorageMapper.to_dict(old_belief)
        statement_data = StatementStorageMapper.to_dict(self.statement)

        def retrieve_by_id_side_effect(collection, doc_id):
            if collection == "statements" and doc_id == self.statement.id:
                return statement_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

        one_minute_ago = datetime.now() - timedelta(minutes=1)
        results = self.repository.get_recently_updated(self.agent.id, one_minute_ago)
        
        self.assertEqual(len(results), 0)

    def test_get_recently_updated_returns_empty_for_no_beliefs(self):
        self.storage_adapter.retrieve_by_attributes.return_value = []
        
        one_minute_ago = datetime.now() - timedelta(minutes=1)
        results = self.repository.get_recently_updated(self.agent.id, one_minute_ago)
        
        expected_query = {'agent_id': "agent_id"}
        self.storage_adapter.retrieve_by_attributes.assert_called_once_with("beliefs", expected_query)
        self.assertEqual(results, [])

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


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