from typing import Dict, List, Optional, Any, TYPE_CHECKING
from datetime import datetime, timedelta

from analogai.foundation.modules.notion.notion_service import NotionService
from analogai.foundation.modules.direct_statement.statement_service_base import StatementServiceBase

if TYPE_CHECKING:
    from analogai.foundation.common.enums.criterion import Criterion
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
from analogai.foundation.modules.belief.belief_assembler import BeliefAssembler
from analogai.foundation.modules.belief.belief_deduction_applier import BeliefDeductionApplier
from analogai.foundation.modules.belief.belief_filterer import BeliefFilterer, BeliefKey
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.common.logging.app_logger import app_logger
from analogai.foundation.common.utils.statement_serializers import serialize_statement
from typeguard import typechecked


class BeliefService:
    def __init__(
        self,
        statement_service: StatementServiceBase,
        notion_service: NotionService,
        categorical_deduction_service=None,
        conditional_deduction_service=None,
    ):
        if statement_service is None:
            raise ValueError("statement_service is required")
        self.statement_service = statement_service
        self.notion_service = notion_service
        self._categorical_deduction_service = categorical_deduction_service
        self._conditional_deduction_service = conditional_deduction_service
        self._belief_assembler = BeliefAssembler()
        self._belief_filterer = BeliefFilterer(
            statement_service=statement_service,
            notion_service=notion_service,
            belief_assembler=self._belief_assembler,
        )
        self._deduction_applier = BeliefDeductionApplier(
            belief_assembler=self._belief_assembler,
            belief_key_builder=self._belief_filterer.belief_key_from_statement,
            categorical_deduction_service=categorical_deduction_service,
            conditional_deduction_service=conditional_deduction_service,
            max_deduction_depth=5,
        )

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

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

    def _base_beliefs_for_agent(
        self,
        agent_id: str,
        modifier: Optional[Modifier],
        quantity: Optional[float],
    ) -> Dict[BeliefKey, Belief]:
        return self._belief_filterer.base_beliefs_for_agent(agent_id, modifier, quantity)

    def _belief_key_from_values(
        self,
        agent_id: str,
        subject_id: str,
        object_id: str,
        relation: Relation,
        modifier: Optional[Modifier],
    ) -> BeliefKey:
        return self._belief_filterer.belief_key_from_values(
            agent_id=agent_id,
            subject_id=subject_id,
            object_id=object_id,
            relation=relation,
            modifier=modifier,
        )

    def _belief_key_from_statement(self, statement: DirectStatement) -> BeliefKey:
        return self._belief_filterer.belief_key_from_statement(statement)

    def _beliefs_for_agent(
        self,
        agent_id: str,
        modifier: Optional[Modifier],
        quantity: Optional[float],
        user_id: Optional[str] = None,
    ) -> Dict[BeliefKey, Belief]:
        beliefs_by_key = self._belief_filterer.base_beliefs_for_agent(agent_id, modifier, quantity)
        return self._deduction_applier.apply_deductions(
            beliefs_by_key,
            agent_id=agent_id,
            modifier=modifier,
            quantity=quantity,
            user_id=user_id,
        )

    def get_base_beliefs_for_agent(
        self,
        agent_id: str,
        modifier: Optional[Modifier] = None,
        quantity: Optional[float] = None,
    ) -> List[Belief]:
        beliefs_by_key = self._belief_filterer.base_beliefs_for_agent(agent_id, modifier, quantity)
        return list(beliefs_by_key.values())


    @typechecked
    def add_statement_to_belief(self, statement: DirectStatement) -> Belief:
        statements = [statement]
        return self._belief_assembler.build_belief(statements)

    @typechecked
    def get_recently_updated(self, agent_id: str, since: Optional[datetime] = None) -> List[Belief]:
        if since is None:
            since = datetime.now() - timedelta(seconds=20)
        beliefs = self.get_all_for_agent(agent_id)
        return [belief for belief in beliefs if belief.updated_at >= since or belief.created_at >= since]

    @typechecked
    def find(self, subject_notion: Notion, complement_notion: Notion,
             relation: Relation,
             modifier: Optional[Modifier] = None,
             quantity: Optional[float] = None,
             user_id: Optional[str] = None) -> Optional[Belief]:
        modifier = self._belief_filterer.apply_quantity_modifier(modifier, quantity)
        belief_key = self._belief_filterer.belief_key_from_values(
            agent_id=subject_notion.agent_id,
            subject_id=subject_notion.id,
            object_id=complement_notion.id,
            relation=relation,
            modifier=modifier,
        )
        beliefs_by_key = self._beliefs_for_agent(
            agent_id=subject_notion.agent_id,
            modifier=modifier,
            quantity=quantity,
            user_id=user_id,
        )
        belief = beliefs_by_key.get(belief_key)
        if belief is None:
            return None
        serialized = serialize_statement(
            subject_notion_caption=belief.subject_notion.caption,
            complement_notion_caption=belief.complement_notion.caption,
            relation=belief.relation,
            probability=float(belief.probability) if belief.probability is not None else 1.0,
            modifier=belief.modifier,
            subject_attributes=belief.subject_notion.attributes,
            complement_attributes=belief.complement_notion.attributes,
        )
        app_logger.debug(f"Found Belief: {serialized} (belief_id: {belief.id})")
        return belief

    @typechecked
    def find_by_expressions(self, agent_id: str, subject_notion_expression: NotionExpression,
                           complement_notion_expression: NotionExpression,
                           relation: Relation,
                           modifier: Optional[Modifier] = None,
                           user_id: Optional[str] = None) -> Optional[Belief]:
        subject_notion = self.notion_service.find(agent_id, subject_notion_expression)
        if not subject_notion:
            app_logger.warning(f"Source notion could not be found using expression '{subject_notion_expression.caption}'")
            return None
        complement_notion = self.notion_service.find(agent_id, complement_notion_expression)
        if not complement_notion:
            app_logger.warning(f"Target notion could not be found using expression '{complement_notion_expression.caption}'")
            return None
        return self.find(
            subject_notion=subject_notion,
            complement_notion=complement_notion,
            relation=relation,
            modifier=modifier,
            user_id=user_id,
        )

    @typechecked
    def search_by_notions(self, subject_notion: Notion, complement_notion: Notion,
                          modifier: Optional[Modifier] = None) -> List[Belief]:
        return self._belief_filterer.search_by_notions(subject_notion, complement_notion, modifier)

    @typechecked
    def search_by_subject_notion(self, subject_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        return self._belief_filterer.search_by_subject_notion(subject_notion, modifier)

    @typechecked
    def search_by_complement_notion(self, complement_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        return self._belief_filterer.search_by_complement_notion(complement_notion, modifier)

    @typechecked
    def search_by_subject_notion_and_relation(self, subject_notion: Notion,
                                                       relation: Relation,
                                                       modifier: Optional[Modifier] = None) -> List[Belief]:
        return self._belief_filterer.search_by_subject_notion_and_relation(subject_notion, relation, modifier)

    @typechecked
    def search_by_complement_notion_and_relation(self, complement_notion: Notion,
                                                      relation: Relation,
                                                      modifier: Optional[Modifier] = None) -> List[Belief]:
        return self._belief_filterer.search_by_complement_notion_and_relation(complement_notion, relation, modifier)

    @typechecked
    def search_by_subject_notion_expression(self, agent_id: str, subject_notion_expression: NotionExpression,
                                            modifier: Optional[Modifier] = None) -> Optional[List[Belief]]:
        return self._belief_filterer.search_by_subject_notion_expression(agent_id, subject_notion_expression, modifier)

    @typechecked
    def search_by_complement_notion_expression(self, agent_id: str, complement_notion_expression: NotionExpression,
                                          modifier: Optional[Modifier] = None) -> Optional[List[Belief]]:
        return self._belief_filterer.search_by_complement_notion_expression(agent_id, complement_notion_expression, modifier)

    @typechecked
    def forget(self, belief: Belief) -> None:
        serialized = serialize_statement(
            subject_notion_caption=belief.subject_notion.caption,
            complement_notion_caption=belief.complement_notion.caption,
            relation=belief.relation,
            probability=float(belief.probability) if belief.probability is not None else 1.0,
            modifier=belief.modifier,
            subject_attributes=belief.subject_notion.attributes,
            complement_attributes=belief.complement_notion.attributes,
        )
        app_logger.debug(f"Skipped persistence for Belief: {serialized} (belief_id: {belief.id})")

    @typechecked
    def get_all_for_agent(
        self,
        agent_id: str,
        user_id: Optional[str] = None,
    ) -> List[Belief]:
        beliefs_by_key = self._beliefs_for_agent(
            agent_id=agent_id,
            modifier=None,
            quantity=None,
            user_id=user_id,
        )
        beliefs = list(beliefs_by_key.values())
        app_logger.debug(f"Found {len(beliefs)} beliefs")
        return beliefs


    @typechecked
    def search_by_criteria(self, agent_id: str, clause: Dict["Criterion", Any]) -> bool:
        from analogai.foundation.common.utils.belief_criteria_matcher import BeliefCriteriaMatcher
        beliefs = self.get_all_for_agent(agent_id)
        return any(BeliefCriteriaMatcher.matches_clause(belief, clause) for belief in beliefs)

    @typechecked
    def search_by_modifier(self, modifier: Modifier) -> List[Belief]:
        app_logger.warning("search_by_modifier is not supported without agent scope in live mode")
        return []