from typing import Optional, Any
from analogai.foundation.extensions.belief_query.services.belief_query_service import BeliefQueryService
from analogai.foundation.common.enums.criterion import Criterion
from analogai.deepthink.application.utils.learnings_search_modifier import apply_modifier
from analogai.foundation.common.utils.expression_factory import ExpressionFactory
from analogai.deepthink.application.usecases.learning.learn_statement.models import LearnStatementRequest, StatementType
from analogai.deepthink.application.usecases.learning.make_direct_statement.models import MakeDirectStatementResponse
from analogai.deepthink.application.value_objects.user_agent import UserAgent


class LearningsServiceAdapter:
    def __init__(
        self,
        belief_query_service: BeliefQueryService,
        notion_service: Any,
        simple_statement_presenter: Any,
    ):
        self._belief_query_service = belief_query_service
        self._notion_service = notion_service
        self._simple_statement_presenter = simple_statement_presenter
        self._expression_factory = ExpressionFactory()

    def learn(self, request: LearnStatementRequest) -> Optional[Any]:
        if request.statement_type == StatementType.SIMPLE:
            return self._learn_simple(request)
        return None

    def _learn_simple(self, request: LearnStatementRequest) -> Optional[Any]:
        if not request.subject or not request.complement or not request.relation:
            return None
        user_id = request.user_agent.user_id
        agent_id = request.user_agent.agent_id
        subject_notion = self._notion_service.find_or_create(
            user_id, agent_id, request.subject
        )
        complement_caption = request.complement.caption
        builder = self._belief_query_service.search(user_id, agent_id).where(
            Criterion.SubjectCaption, subject_notion.caption
        ).where(Criterion.ComplementCaption, complement_caption).where(
            Criterion.Relation, request.relation
        )
        apply_modifier(builder, request.modifier)
        beliefs = builder.list()
        if beliefs:
            belief = beliefs[0]
            response = MakeDirectStatementResponse(
                user_agent=request.user_agent,
                belief_expression=self._expression_factory.propagate_belief_expression(belief),
            )
            return self._simple_statement_presenter.output_statement_already_known(response)

        result_criteria = {
            Criterion.Relation: request.relation,
            Criterion.ComplementCaption: complement_caption,
        }
        learning = self._belief_query_service.create(
            user_id=user_id,
            agent_id=agent_id,
            subject_notion=subject_notion,
            result_criteria=result_criteria,
            modifier=request.modifier,
            probability=request.probability if request.probability is not None else 1.0,
        )
        belief = self._belief_query_service.resolve_to_entity(learning)
        if not belief:
            return None
        response = MakeDirectStatementResponse(
            user_agent=request.user_agent,
            belief_expression=self._expression_factory.propagate_belief_expression(belief),
        )
        return self._simple_statement_presenter.output_create_statement_success(response)

