from typing import Dict, List, Optional, Any, TYPE_CHECKING, Iterable, Tuple
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.shared.statement_base import StatementBase
from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.core.value_objects.time import Time
from analogai.foundation.core.value_objects.attribute import Attribute
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.common.logging.app_logger import app_logger
from analogai.foundation.common.utils.modifier_utils.modifier_matcher import ModifierMatcher
from analogai.foundation.common.utils.statement_serializers import serialize_statement
from typeguard import typechecked


BeliefKey = Tuple[str, str, Relation, str, Optional[Tuple]]


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._matcher = ModifierMatcher()
        self._max_deduction_depth = 5

    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 _time_key(self, time: Optional[Time]) -> Optional[Tuple[Optional[int], Optional[int], Optional[int], Optional[int], Optional[int]]]:
        if time is None:
            return None
        return (time.year, time.month, time.day, time.hour, time.minute)

    def _attributes_key(self, attributes: Optional[List[Attribute]]) -> Optional[Tuple[Tuple[Optional[str], Optional[float], Optional[str]], ...]]:
        if attributes is None:
            return None
        return tuple((attribute.tag, attribute.value, attribute.unit) for attribute in attributes)

    def _modifier_has_context(self, modifier: Optional[Modifier]) -> bool:
        if modifier is None:
            return False
        if modifier.location is not None:
            return True
        if modifier.from_time is not None:
            return True
        if modifier.to_time is not None:
            return True
        if modifier.deontic_modality is not None:
            return True
        if modifier.attributes:
            return True
        if modifier.quantity is not None:
            return True
        return False

    def _modifier_key(self, modifier: Optional[Modifier]) -> Optional[Tuple]:
        if not self._modifier_has_context(modifier):
            return None
        assert modifier is not None
        deontic = modifier.deontic_modality.value if modifier.deontic_modality is not None else None
        return (
            modifier.location,
            deontic,
            self._time_key(modifier.from_time),
            self._time_key(modifier.to_time),
            self._attributes_key(modifier.attributes),
            modifier.quantity,
        )

    def _apply_quantity_modifier(self, modifier: Optional[Modifier], quantity: Optional[float]) -> Optional[Modifier]:
        if quantity is None:
            return modifier
        if modifier is None:
            return Modifier(quantity=quantity)
        return Modifier(
            location=modifier.location,
            deontic_modality=modifier.deontic_modality,
            from_time=modifier.from_time,
            to_time=modifier.to_time,
            attributes=list(modifier.attributes) if modifier.attributes else None,
            quantity=quantity,
        )

    def _belief_key_from_statement(self, statement: StatementBase) -> BeliefKey:
        return self._belief_key_from_values(
            agent_id=statement.agent_id,
            subject_id=statement.subject_notion.id,
            object_id=statement.complement_notion.id,
            relation=statement.relation,
            modifier=statement.modifier,
        )

    def _belief_key_from_values(
        self,
        agent_id: str,
        subject_id: str,
        object_id: str,
        relation: Relation,
        modifier: Optional[Modifier],
    ) -> BeliefKey:
        if not isinstance(relation, Relation):
            raise TypeError("relation must be a Relation instance")
        return (
            agent_id,
            subject_id,
            relation,
            object_id,
            self._modifier_key(modifier),
        )

    def _modifier_matches(self, stored: Optional[Modifier], query: Optional[Modifier]) -> bool:
        stored_normalized = stored if self._modifier_has_context(stored) else None
        query_normalized = query if self._modifier_has_context(query) else None
        return self._matcher.context_matches(stored_normalized, query_normalized)

    def _modifier_exact_matches(self, stored: Optional[Modifier], query: Optional[Modifier]) -> bool:
        stored_normalized = stored if self._modifier_has_context(stored) else None
        query_normalized = query if self._modifier_has_context(query) else None
        return self._matcher.context_exact_matches(stored_normalized, query_normalized)

    def _resolve_notion_by_id(self, notion_id: str, agent_id: str) -> Notion:
        repository = getattr(self.notion_service, "notion_repository", None)
        if repository is not None:
            found = repository.find_by_id(notion_id)
            if found:
                return found
        return Notion(user_id="system", agent_id=agent_id, caption=notion_id, id=notion_id)

    def _build_belief(self, statements: Iterable[StatementBase], belief_id: Optional[str] = None) -> Belief:
        statements_list = list(statements)
        if not statements_list:
            raise ValueError("Cannot build belief without statements")
        ordered = sorted(
            statements_list,
            key=lambda stmt: getattr(stmt, "created_at", datetime.min),
        )
        first = ordered[0]
        belief = Belief(
            agent_id=first.agent_id,
            id=belief_id or first.id,
        )
        for statement in ordered:
            belief.add_statement(statement)
        created_candidates: List[datetime] = []
        for stmt in ordered:
            created_at = getattr(stmt, "created_at", None)
            if isinstance(created_at, datetime):
                created_candidates.append(created_at)
        created_at = min(created_candidates) if created_candidates else None
        updated_at = max(created_candidates) if created_candidates else None
        if created_at:
            belief.created_at = created_at
        if updated_at:
            belief.updated_at = updated_at
        return belief

    def _group_statements(self, statements: Iterable[StatementBase]) -> Dict[BeliefKey, List[StatementBase]]:
        groups: Dict[BeliefKey, List[StatementBase]] = {}
        for statement in statements:
            key = self._belief_key_from_statement(statement)
            groups.setdefault(key, []).append(statement)
        return groups

    def _base_beliefs_for_agent(
        self,
        agent_id: str,
        modifier: Optional[Modifier],
        quantity: Optional[float],
    ) -> Dict[BeliefKey, Belief]:
        modifier = self._apply_quantity_modifier(modifier, quantity)
        statements = self.statement_service.find_by_agent_id(agent_id)
        if modifier is None:
            filtered = list(statements)
        else:
            filtered = [
                stmt for stmt in statements
                if self._modifier_exact_matches(stmt.modifier, modifier)
            ]
        groups = self._group_statements(filtered)
        return {
            belief_key: self._build_belief(group)
            for belief_key, group in groups.items()
        }

    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_from_statement(deduction)
                belief = beliefs_by_key.get(belief_key)
                if belief is None:
                    beliefs_by_key[belief_key] = self._build_belief([deduction])
                else:
                    belief.add_statement(deduction)

        return beliefs_by_key

    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._base_beliefs_for_agent(agent_id, modifier, quantity)
        self._apply_deductions(
            beliefs_by_key,
            agent_id=agent_id,
            modifier=modifier,
            quantity=quantity,
            user_id=user_id,
        )
        return beliefs_by_key

    def get_base_beliefs_for_agent(
        self,
        agent_id: str,
        modifier: Optional[Modifier] = None,
        quantity: Optional[float] = None,
    ) -> List[Belief]:
        beliefs_by_key = self._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._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._apply_quantity_modifier(modifier, quantity)
        belief_key = self._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]:
        statements = self.statement_service.find_by_subject(subject_notion)
        filtered = [
            stmt for stmt in statements
            if stmt.complement_notion.id == complement_notion.id
            and self._modifier_matches(stmt.modifier, modifier)
        ]
        groups = self._group_statements(filtered)
        beliefs = [self._build_belief(group) for group in groups.values()]
        beliefs.sort(key=lambda belief: (belief.updated_at, belief.id), reverse=True)
        app_logger.debug(f"Found {len(beliefs)} beliefs - subject='{subject_notion.caption}', object='{complement_notion.caption}'")
        return beliefs

    @typechecked
    def search_by_subject_notion(self, subject_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        statements = self.statement_service.find_by_subject(subject_notion)
        filtered = [stmt for stmt in statements if self._modifier_matches(stmt.modifier, modifier)]
        groups = self._group_statements(filtered)
        beliefs = [self._build_belief(group) for group in groups.values()]
        beliefs.sort(key=lambda belief: (belief.updated_at, belief.id), reverse=True)
        app_logger.debug(f"Found {len(beliefs)} beliefs - subject='{subject_notion.caption}'")
        return beliefs

    @typechecked
    def search_by_complement_notion(self, complement_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        statements = self.statement_service.find_by_object(complement_notion)
        filtered = [stmt for stmt in statements if self._modifier_matches(stmt.modifier, modifier)]
        groups = self._group_statements(filtered)
        beliefs = [self._build_belief(group) for group in groups.values()]
        beliefs.sort(key=lambda belief: (belief.updated_at, belief.id), reverse=True)
        app_logger.debug(f"Found {len(beliefs)} beliefs - object='{complement_notion.caption}'")
        return beliefs

    @typechecked
    def search_by_subject_notion_and_relation(self, subject_notion: Notion,
                                                       relation: Relation,
                                                       modifier: Optional[Modifier] = None) -> List[Belief]:
        beliefs = self.search_by_subject_notion(subject_notion=subject_notion, modifier=modifier)
        return [b for b in beliefs if b.relation == relation]

    @typechecked
    def search_by_complement_notion_and_relation(self, complement_notion: Notion,
                                                      relation: Relation,
                                                      modifier: Optional[Modifier] = None) -> List[Belief]:
        beliefs = self.search_by_complement_notion(complement_notion=complement_notion, modifier=modifier)
        return [b for b in beliefs if b.relation == relation]

    @typechecked
    def search_by_subject_notion_expression(self, agent_id: str, subject_notion_expression: NotionExpression,
                                            modifier: Optional[Modifier] = None) -> Optional[List[Belief]]:
        subject_notion = self.notion_service.find(agent_id, subject_notion_expression)
        if not subject_notion:
            return None
        return self.search_by_subject_notion(subject_notion=subject_notion, modifier=modifier)

    @typechecked
    def search_by_complement_notion_expression(self, agent_id: str, complement_notion_expression: NotionExpression,
                                          modifier: Optional[Modifier] = None) -> Optional[List[Belief]]:
        complement_notion = self.notion_service.find(agent_id, complement_notion_expression)
        if not complement_notion:
            return None
        return self.search_by_complement_notion(complement_notion=complement_notion, modifier=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.extensions.belief_query.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 []