from typing import Optional, Dict, Iterable, List, Tuple, TYPE_CHECKING
from datetime import datetime

from analogai.foundation.common.logging.app_logger import app_logger
from analogai.foundation.common.utils.statement_serializers import serialize_statement
from analogai.foundation.common.utils.relationship_type_storage import select_relation
from analogai.foundation.common.utils.modifier_utils.modifier_matcher import ModifierMatcher
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.categorical_deduction.categorical_deduction import CategoricalDeduction
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.shared.statement_base import StatementBase
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
if TYPE_CHECKING:
    from analogai.foundation.modules.belief.belief_service import BeliefService


class CategoricalDeductionService:
    def __init__(
        self,
        belief_service: "BeliefService",
    ):
        if belief_service is None:
            raise ValueError("BeliefService is required for categorical deductions")
        self._belief_service = belief_service
        self._matcher = ModifierMatcher()

    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 _belief_key_from_statement(self, statement: StatementBase) -> Tuple:
        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],
    ) -> Tuple:
        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_exact_matches(self, stored: Optional[Modifier], query: Optional[Modifier]) -> bool:
        return self._matcher.context_exact_matches(stored, query)

    def _belief_matches_query(self, belief: Belief, modifier: Optional[Modifier]) -> bool:
        return self._modifier_exact_matches(belief.modifier, modifier)

    def _index_beliefs_by_subject(self, beliefs: Iterable[Belief]) -> Dict[str, List[Belief]]:
        index: Dict[str, List[Belief]] = {}
        for belief in beliefs:
            index.setdefault(belief.subject_notion.id, []).append(belief)
        return index

    def _build_belief(self, statements: List[StatementBase]) -> Belief:
        if not statements:
            raise ValueError("Cannot build belief without statements")
        ordered = sorted(
            statements,
            key=lambda stmt: getattr(stmt, "created_at", datetime.min),
        )
        first = ordered[0]
        belief = Belief(
            agent_id=first.agent_id,
            id=first.id,
        )
        for statement in ordered:
            belief.add_statement(statement)
        return belief

    def find_deductions_by_agent_id(
        self,
        agent_id: str,
        modifier: Optional[Modifier] = None,
        max_depth: int = 5,
        limit: Optional[int] = None,
        offset: int = 0,
    ) -> List[StatementBase]:
        if self._belief_service is None:
            raise ValueError("BeliefService is not configured")
        getter = getattr(self._belief_service, "get_base_beliefs_for_agent", None)
        if callable(getter):
            try:
                beliefs = getter(agent_id, modifier=modifier, quantity=None)
            except TypeError:
                beliefs = getter(agent_id)
        else:
            beliefs = self._belief_service.get_all_for_agent(agent_id)
        filtered = [
            belief for belief in beliefs
            if self._belief_matches_query(belief, modifier)
        ]
        if offset or limit is not None:
            start = offset
            end = offset + limit if limit is not None else None
            filtered = filtered[start:end]
        beliefs_by_key = {
            self._belief_key_from_values(
                agent_id=belief.agent_id,
                subject_id=belief.subject_notion.id,
                object_id=belief.complement_notion.id,
                relation=belief.relation,
                modifier=belief.modifier,
            ): belief
            for belief in filtered
        }
        return self.resolve_deductions(
            beliefs_by_key=beliefs_by_key,
            modifier=modifier,
            max_depth=max_depth,
        )

    def resolve_deductions(
        self,
        beliefs_by_key: Dict[Tuple, Belief],
        modifier: Optional[Modifier],
        max_depth: int,
    ) -> List[StatementBase]:
        deductions: List[StatementBase] = []
        seen: set[tuple] = set()

        for depth in range(1, max_depth + 1):
            modifier_key = self._modifier_key(modifier)
            # app_logger.info(
            #     "Categorical deduction depth=%s: beliefs=%s, modifier=%s",
            #     depth,
            #     len(beliefs_by_key),
            #     modifier_key,
            # )
            index = self._index_beliefs_by_subject(beliefs_by_key.values())
            new_added = False
            for minor in list(beliefs_by_key.values()):
                if not self._belief_matches_query(minor, modifier):
                    continue
                major_candidates = index.get(minor.complement_notion.id, [])
                for major in major_candidates:
                    if not self._belief_matches_query(major, modifier):
                        continue
                    relation = select_relation(major.relation, minor.relation)
                    if relation is None:
                        continue
                    if minor.subject_notion.id == major.complement_notion.id:
                        continue
                    key = (
                        minor.subject_notion.id,
                        major.complement_notion.id,
                        relation,
                        major.id,
                        minor.id,
                        modifier_key,
                    )
                    if key in seen:
                        continue
                    seen.add(key)
                    # app_logger.info(
                    #     "Categorical deduction depth=%s: minor=%s -> major=%s => relation=%s",
                    #     depth,
                    #     minor.id,
                    #     major.id,
                    #     relation,
                    # )
                    deduction = self.resolve(
                        agent_id=minor.agent_id,
                        major_premise=major,
                        minor_premise=minor,
                        subject_notion=minor.subject_notion,
                        complement_notion=major.complement_notion,
                        relation=relation,
                        modifier=modifier,
                        deduction_id=None,
                    )
                    deductions.append(deduction)
                    belief_key = self._belief_key_from_statement(deduction)
                    if belief_key in beliefs_by_key:
                        beliefs_by_key[belief_key].add_statement(deduction)
                    else:
                        beliefs_by_key[belief_key] = self._build_belief([deduction])
                    new_added = True
            if not new_added:
                # app_logger.info("Categorical deduction stopped at depth=%s: no new deductions", depth)
                break
        return deductions

    def resolve(
        self,
        agent_id: str,
        major_premise: Belief,
        minor_premise: Belief,
        subject_notion: Notion,
        complement_notion: Notion,
        relation: Relation,
        modifier: Optional[Modifier] = None,
        deduction_id: Optional[str] = None,
    ) -> CategoricalDeduction:
        if subject_notion.id == complement_notion.id:
            raise ValueError(
                f"Cannot resolve categorical deduction: source notion '{subject_notion.caption}' (id={subject_notion.id}) "
                "cannot be the same as target notion"
            )

        deduction = CategoricalDeduction(
            agent_id=agent_id,
            major_premise=major_premise,
            minor_premise=minor_premise,
            subject_notion=subject_notion,
            complement_notion=complement_notion,
            relation=relation,
            modifier=modifier,
            id=deduction_id,
        )

        major_serialized = serialize_statement(
            subject_notion_caption=major_premise.subject_notion.caption,
            complement_notion_caption=major_premise.complement_notion.caption,
            relation=major_premise.relation,
            probability=major_premise.probability,
            modifier=major_premise.modifier,
            subject_attributes=major_premise.subject_notion.attributes,
            complement_attributes=major_premise.complement_notion.attributes,
        )
        minor_serialized = serialize_statement(
            subject_notion_caption=minor_premise.subject_notion.caption,
            complement_notion_caption=minor_premise.complement_notion.caption,
            relation=minor_premise.relation,
            probability=minor_premise.probability,
            modifier=minor_premise.modifier,
            subject_attributes=minor_premise.subject_notion.attributes,
            complement_attributes=minor_premise.complement_notion.attributes,
        )
        app_logger.debug(f"Calculated probability: {deduction.probability} and validity: {deduction.validity} from:")
        app_logger.debug(
            f"Major premise: {major_serialized} (prob={major_premise.probability}, valid={major_premise.validity})"
        )
        app_logger.debug(
            f"Minor premise: {minor_serialized} (prob={minor_premise.probability}, valid={minor_premise.validity})"
        )

        deduction_serialized = serialize_statement(
            subject_notion_caption=deduction.subject_notion.caption,
            complement_notion_caption=deduction.complement_notion.caption,
            relation=deduction.relation,
            probability=float(deduction.probability) if deduction.probability is not None else 1.0,
            modifier=deduction.modifier,
            subject_attributes=deduction.subject_notion.attributes,
            complement_attributes=deduction.complement_notion.attributes,
        )
        app_logger.debug(f"Resolved CategoricalDeduction: {deduction_serialized} (id: {deduction.id})")

        return deduction