from unittest.mock import Mock
from analogai.foundation.common.enums import RelationshipType
from analogai.foundation.common.dtos.belief_expression import Relation
from analogai.deepthink.startup.factories.domain_service_factory import DomainServiceFactory
from analogai.deepthink.application.usecases.learning.make_direct_statement.models import MakeDirectStatementRequest
from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.modules.user.user import User
from analogai.foundation.modules.agent.agent import Agent
from analogai.deepthink.tests.functional.storage_cleaner import cleanup_storages
from analogai.deepthink.tests.mocks.output_ports.mock_learning_output_port import MockLearningOutputPort


class LearnerTestSetup:
    def __init__(self):
        from analogai.foundation.settings import config as foundation_config

        self._config_override = None
        self._config_override = foundation_config.override(
            **{"vector.enabled": False}
        )
        self._config_override.__enter__()

        domain_service_factory = DomainServiceFactory()

        self.notion_service = domain_service_factory.get_notion_service()
        self.belief_service = domain_service_factory.get_belief_service()
        self.statement_service = domain_service_factory.get_statement_service()
        self.statement_service = self.statement_service

        self.user = User(id="test-user", name="Test User", email="test@example.com")
        self.agent = Agent(id="test-agent", knowledge_base="Test KB", creator=self.user)
        self.user_agent = UserAgent(
            user_id=self.user.id,
            agent_id=self.agent.id,
            user_authority=1.0
        )
        self.mock_output_port = MockLearningOutputPort()
        self.relation = Relation.from_type(RelationshipType.RELATED_TO)
        self.expected_user = self.user

    def create_learner_request(self, probability: float = 0.1, subject_expression=None, object_expression=None, relation=None, modifier=None):
        return MakeDirectStatementRequest(
            user_agent=self.user_agent,
            subject_notion_expression=subject_expression or NotionExpression(caption="Source"),
            complement_notion_expression=object_expression or NotionExpression(caption="Target"),
            relation=relation or self.relation,
            probability=probability,
            modifier=modifier,
        )

    def clear(self):
        if self._config_override is not None:
            self._config_override.__exit__(None, None, None)
            self._config_override = None
        cleanup_storages()


