from typing import List, Dict, Any, Optional, Protocol, runtime_checkable

from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.extensions.belief_query.services.belief_query_service import BeliefQueryService


class TraversableNotion:
    def __init__(self, notion: Notion, learnings_service: BeliefQueryService):
        self._notion = notion
        self._learnings_service = learnings_service

    def query_in(self, **criteria: Any) -> List[Belief]:
        builder = self._learnings_service.search(
            self._notion.user_id, self._notion.agent_id
        ).where(Criterion.ComplementCaption, self._notion.caption)
        if criteria:
            criterion_map = self._to_criterion_map(criteria)
            builder = builder.apply(criterion_map)
        return builder.list()

    def get_learnings_out(self, **criteria: Any) -> List[Belief]:
        builder = self._learnings_service.search(
            self._notion.user_id, self._notion.agent_id
        ).where(Criterion.SubjectCaption, self._notion.caption)
        if criteria:
            criterion_map = self._to_criterion_map(criteria)
            builder = builder.apply(criterion_map)
        return builder.list()

    def get_beliefs_in(self, relation: Relation, modifier: Optional[Modifier] = None) -> List[Belief]:
        criteria = self._build_belief_criteria(relation, modifier)
        builder = self._learnings_service.search(
            self._notion.user_id, self._notion.agent_id
        ).where(Criterion.ComplementCaption, self._notion.caption)
        if criteria:
            builder = builder.apply(criteria)
        return builder.list()

    def query_out(self, relation: Relation, modifier: Optional[Modifier] = None) -> List[Belief]:
        criteria = self._build_belief_criteria(relation, modifier)
        builder = self._learnings_service.search(
            self._notion.user_id, self._notion.agent_id
        ).where(Criterion.SubjectCaption, self._notion.caption)
        if criteria:
            builder = builder.apply(criteria)
        return builder.list()

    def _build_belief_criteria(
        self, relation: Relation, modifier: Optional[Modifier]
    ) -> Dict[Criterion, Any]:
        if not isinstance(relation, Relation):
            raise TypeError("relation must be a Relation instance")
        criteria: Dict[Criterion, Any] = {
            Criterion.Relation: relation,
        }
        if modifier is None:
            return criteria
        if modifier.location is not None:
            criteria[Criterion.Location] = modifier.location
        if modifier.deontic_modality is not None:
            criteria[Criterion.DeonticModality] = modifier.deontic_modality
        if modifier.from_time is not None:
            criteria[Criterion.FromTime] = modifier.from_time
        if modifier.to_time is not None:
            criteria[Criterion.ToTime] = modifier.to_time
        if modifier.attributes:
            criteria[Criterion.AttributeValue] = modifier.attributes
        return criteria

    def _to_criterion_map(self, criteria: Dict[str, Any]) -> Dict[Criterion, Any]:
        result = {}
        for k, v in criteria.items():
            try:
                result[Criterion(k)] = v
            except ValueError:
                pass
        return result

    @property
    def id(self) -> str:
        return self._notion.id

    @property
    def caption(self) -> str:
        return self._notion.caption

    @property
    def user_id(self) -> str:
        return self._notion.user_id

    @property
    def agent_id(self) -> str:
        return self._notion.agent_id

    @property
    def notion(self) -> Notion:
        return self._notion


def get_traversable_notion(
    belief_query_service: BeliefQueryService,
    notion: Notion,
) -> TraversableNotion:
    if isinstance(belief_query_service, TraversableNotionProvider):
        traversable = belief_query_service.get_traversable_notion(notion)
        if traversable is not None:
            return traversable
    return TraversableNotion(notion, belief_query_service)


@runtime_checkable
class TraversableNotionProvider(Protocol):
    def get_traversable_notion(self, notion: Notion) -> Optional[TraversableNotion]:
        ...
