from typing import Dict, Optional, Callable

from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.shared.statement_base import StatementBase
from analogai.foundation.modules.belief.belief_filterer import BeliefKey
from analogai.foundation.core.value_objects.modifier import Modifier


class BeliefDeductionApplier:
    def __init__(
        self,
        belief_assembler,
        belief_key_builder: Callable[[StatementBase], BeliefKey],
        categorical_deduction_service=None,
        conditional_deduction_service=None,
        max_deduction_depth: int = 5,
    ):
        self._belief_assembler = belief_assembler
        self._belief_key_builder = belief_key_builder
        self._categorical_deduction_service = categorical_deduction_service
        self._conditional_deduction_service = conditional_deduction_service
        self._max_deduction_depth = max_deduction_depth

    def set_categorical_deduction_service(self, service) -> None:
        self._categorical_deduction_service = service

    def set_conditional_deduction_service(self, service) -> None:
        self._conditional_deduction_service = service

    def apply_deductions(
        self,
        beliefs_by_key: Dict[BeliefKey, Belief],
        agent_id: str,
        modifier: Optional[Modifier],
        quantity: Optional[float],
        user_id: Optional[str] = None,
    ) -> Dict[BeliefKey, Belief]:
        if self._categorical_deduction_service is not None:
            self._categorical_deduction_service.resolve_deductions(
                beliefs_by_key,
                modifier=modifier,
                max_depth=self._max_deduction_depth,
            )

        if self._conditional_deduction_service is not None:
            deductions = self._conditional_deduction_service.resolve_deductions(
                agent_id=agent_id,
                beliefs=list(beliefs_by_key.values()),
                modifier=modifier,
                user_id=user_id,
            )
            for deduction in deductions:
                belief_key = self._belief_key_builder(deduction)
                belief = beliefs_by_key.get(belief_key)
                if belief is None:
                    beliefs_by_key[belief_key] = self._belief_assembler.build_belief([deduction])
                else:
                    belief.add_statement(deduction)

        return beliefs_by_key
