from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple, Union

from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.belief.belief_service import BeliefService
from analogai.foundation.extensions.belief_query.condition import Condition
from analogai.foundation.extensions.belief_query.belief_criteria_matcher import BeliefCriteriaMatcher
from analogai.foundation.common.utils.similarity.similarity_service import SimilarityService
from analogai.foundation.common.dtos.notion_expression import NotionExpression
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.common.enums.constraint_operator import ConstraintOperator
from analogai.foundation import config

CriterionConstraint = Union[Any, Tuple[Union[str, ConstraintOperator], Any], Condition]


class LearningsSearchBuilder:
    def __init__(
        self,
        belief_service: BeliefService,
        user_id: str,
        agent_id: str,
        similarity_service: Optional[SimilarityService] = None,
    ):
        self._belief_service = belief_service
        self._user_id = user_id
        self._agent_id = agent_id
        self._similarity_service = similarity_service
        self._modifier_strategy = config.BELIEF_QUERY_MODIFIER_STRATEGY
        self._main_clause: Dict[Criterion, Union[CriterionConstraint, List[CriterionConstraint]]] = {}
        self._or_clauses: List[Dict[Criterion, CriterionConstraint]] = []

    def exactModifierMatch(self) -> "LearningsSearchBuilder":
        self._modifier_strategy = "exact"
        return self

    def ignoreModifier(self) -> "LearningsSearchBuilder":
        self._modifier_strategy = "ignore"
        return self

    def _add_constraint(self, criterion: Criterion, constraint: CriterionConstraint) -> None:
        existing = self._main_clause.get(criterion)
        if existing is None:
            self._main_clause[criterion] = constraint
            return
        if type(existing) is list:
            existing.append(constraint)
        else:
            self._main_clause[criterion] = [existing, constraint]

    def where(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._main_clause[criterion] = value
        return self

    def whereMin(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._add_constraint(criterion, (ConstraintOperator.Min, value))
        return self

    def whereMax(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._add_constraint(criterion, (ConstraintOperator.Max, value))
        return self

    def whereNotNull(self, criterion: Criterion) -> "LearningsSearchBuilder":
        self._add_constraint(criterion, (ConstraintOperator.NotNull,))
        return self

    def whereNull(self, criterion: Criterion) -> "LearningsSearchBuilder":
        self._add_constraint(criterion, (ConstraintOperator.Null,))
        return self

    def whereGroup(self, criteria: Dict[Criterion, Any]) -> "LearningsSearchBuilder":
        self._main_clause = dict(criteria)
        return self

    def orWhere(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: value})
        return self

    def orWhereMin(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: (ConstraintOperator.Min, value)})
        return self

    def orWhereMax(self, criterion: Criterion, value: Any) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: (ConstraintOperator.Max, value)})
        return self

    def orWhereNotNull(self, criterion: Criterion) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: (ConstraintOperator.NotNull,)})
        return self

    def orWhereNull(self, criterion: Criterion) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: (ConstraintOperator.Null,)})
        return self

    def orWhereGroup(self, criteria: Dict[Criterion, Any]) -> "LearningsSearchBuilder":
        self._or_clauses.append(dict(criteria))
        return self

    def whereIn(self, criterion: Criterion, values: List[Any]) -> "LearningsSearchBuilder":
        self._main_clause[criterion] = list(values)
        return self

    def orWhereIn(self, criterion: Criterion, values: List[Any]) -> "LearningsSearchBuilder":
        self._or_clauses.append({criterion: list(values)})
        return self

    def apply(self, criteria: Dict[Criterion, Any]) -> "LearningsSearchBuilder":
        for criterion, value in criteria.items():
            self._main_clause[criterion] = value
        return self

    def exists(self) -> bool:
        return len(self._execute()) > 0

    def list(self) -> List[Belief]:
        return self._execute()

    def count(self) -> int:
        return len(self._execute())

    def _execute(self) -> List[Belief]:
        beliefs = self._prefetch_beliefs()
        confirmed = [b for b in beliefs if self._is_confirmed_learning(b)]
        return [belief for belief in confirmed if self._matches_criteria(belief)]

    def _prefetch_beliefs(self) -> List[Belief]:
        if self._or_clauses:
            return self._belief_service.get_all_for_agent(
                self._agent_id,
                user_id=self._user_id,
            )

        if self._deductions_enabled():
            return self._belief_service.get_all_for_agent(
                self._agent_id,
                user_id=self._user_id,
            )

        if self._has_multi_value_constraint(Criterion.SubjectCaption) or self._has_multi_value_constraint(
            Criterion.ComplementCaption
        ) or self._has_multi_value_constraint(Criterion.ObjectCaption):
            return self._belief_service.get_all_for_agent(
                self._agent_id,
                user_id=self._user_id,
            )

        subject_caption = self._extract_caption_for_criterion(Criterion.SubjectCaption)
        complement_caption = self._extract_caption_for_criterion(
            Criterion.ComplementCaption
        ) or self._extract_caption_for_criterion(Criterion.ObjectCaption)

        if subject_caption:
            beliefs = self._belief_service.search_by_subject_notion_expression(
                agent_id=self._agent_id,
                subject_notion_expression=NotionExpression(str(subject_caption)),
            )
            return beliefs or []

        if complement_caption:
            beliefs = self._belief_service.search_by_complement_notion_expression(
                agent_id=self._agent_id,
                complement_notion_expression=NotionExpression(str(complement_caption)),
            )
            return beliefs or []

        return self._belief_service.get_all_for_agent(
            self._agent_id,
            user_id=self._user_id,
        )

    def _has_multi_value_constraint(self, criterion: Criterion) -> bool:
        constraint = self._main_clause.get(criterion)
        if isinstance(constraint, list):
            return True
        return False

    def _deductions_enabled(self) -> bool:
        return bool(
            getattr(self._belief_service, "_categorical_deduction_service", None)
            or getattr(self._belief_service, "_conditional_deduction_service", None)
        )

    def _extract_caption_for_criterion(self, criterion: Criterion) -> Optional[Any]:
        if criterion in self._main_clause:
            value = self._extract_constraint_value(self._main_clause.get(criterion))
            if value is not None:
                return value
        for clause in self._or_clauses:
            if criterion in clause:
                value = self._extract_constraint_value(clause.get(criterion))
                if value is not None:
                    return value
        return None

    def _extract_constraint_value(self, constraint: Any) -> Optional[Any]:
        if constraint is None:
            return None
        if isinstance(constraint, list):
            for item in constraint:
                value = self._extract_constraint_value(item)
                if value is not None:
                    return value
            return None
        if isinstance(constraint, Condition):
            return constraint.value
        if isinstance(constraint, (tuple, list)):
            if len(constraint) > 1:
                return constraint[1]
            return None
        return constraint

    def _is_confirmed_learning(self, belief: Belief) -> bool:
        if not belief.statements:
            return False
        return len(belief.statements) > 0

    def _matches_criteria(self, belief: Belief) -> bool:
        if not self._main_clause and not self._or_clauses:
            return True
        main_clause = self._effective_clause(self._main_clause)
        if self._main_clause and not main_clause:
            return True
        if main_clause and self._matches_clause(belief, main_clause):
            return True
        for clause in self._or_clauses:
            effective = self._effective_clause(clause)
            if clause and not effective:
                return True
            if effective and self._matches_clause(belief, effective):
                return True
        return False

    def _effective_clause(self, clause: Dict[Criterion, Any]) -> Dict[Criterion, Any]:
        if self._modifier_strategy != "ignore":
            return clause
        modifier_criteria = {
            Criterion.Location,
            Criterion.Time,
            Criterion.DeonticModality,
            Criterion.Quantity,
        }
        return {
            criterion: constraint
            for criterion, constraint in clause.items()
            if criterion not in modifier_criteria
        }

    def _matches_clause(
        self,
        belief: Belief,
        clause: Dict[Criterion, Union[CriterionConstraint, List[CriterionConstraint]]],
    ) -> bool:
        return BeliefCriteriaMatcher.matches_clause(
            belief,
            clause,
            similarity_service=self._similarity_service,
        )
