from typing import Optional, Dict, Any
from analogai.foundation import config
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.categorical_deduction.categorical_deduction import CategoricalDeduction
from analogai.foundation.common.dtos.notion_expression import NotionExpression


class ExpressionFactory:
    def __init__(self):
        pass

    def _maybe_lemmatize_text(self, text: str) -> str:
        if not config.LEMMATIZATION_ENABLED:
            return text
        from analogai.foundation.common.utils.lemmatization import lemmatize_text
        return lemmatize_text(text)

    def _lemmatize_criteria(self, criteria: Dict[Criterion, Any]) -> Dict[Criterion, Any]:
        if not criteria:
            return {}
        normalized = dict(criteria)
        subject = normalized.get(Criterion.SubjectCaption)
        if subject is not None:
            normalized[Criterion.SubjectCaption] = self._maybe_lemmatize_text(str(subject))
        obj = normalized.get(Criterion.ObjectCaption)
        if obj is not None:
            normalized[Criterion.ObjectCaption] = self._maybe_lemmatize_text(str(obj))
        return normalized
    
    def create_notion_expression(self, caption: str, attributes=None) -> NotionExpression:
        lemmatized_caption = self._maybe_lemmatize_text(caption)
        return NotionExpression(
            caption=lemmatized_caption,
            children=[],
            context=[],
            attributes=attributes,
        )
    
    def propagate_notion_expression(self, notion: Notion) -> NotionExpression:
        lemmatized_caption = self._maybe_lemmatize_text(notion.caption)
        return NotionExpression(
            caption=lemmatized_caption,
            children=[],
            context=[],
            attributes=notion.attributes,
        )
    
    def propagate_belief_expression(self, belief: Belief) -> "BeliefExpression":
        from analogai.foundation.common.dtos.belief_expression import BeliefExpression
        subject_expression = self.propagate_notion_expression(belief.subject_notion)
        object_expression = self.propagate_notion_expression(belief.complement_notion)
        
        return BeliefExpression(
            subject_notion_expression=subject_expression,
            complement_notion_expression=object_expression,
            relation=belief.relation,
            probability=belief.probability,
            validity=belief.validity,
            modifier=belief.modifier,
        )
    
    def propagate_categorical_deduction_expression(self, inference: CategoricalDeduction) -> "CategoricalDeductionExpression":
        from analogai.foundation.common.dtos.categorical_deduction_expression import CategoricalDeductionExpression
        major_premise_expression = self.propagate_belief_expression(inference.major_premise)
        minor_premise_expression = self.propagate_belief_expression(inference.minor_premise)
        subject_expression = self.propagate_notion_expression(inference.subject_notion)
        object_expression = self.propagate_notion_expression(inference.complement_notion)
        
        return CategoricalDeductionExpression(
            major_premise_expression=major_premise_expression,
            minor_premise_expression=minor_premise_expression,
            subject_notion_expression=subject_expression,
            complement_notion_expression=object_expression,
            relation=inference.relation,
            probability=inference.probability,
            validity=inference.validity,
            modifier=inference.modifier
        )

    def propagate_hypothetical_statement_expression(self, hypothetical: "HypotheticalStatement") -> "HypotheticalStatementExpression":
        from analogai.foundation.common.dtos.hypothetical_statement_expression import HypotheticalStatementExpression
        from analogai.foundation.modules.conditionals.hypothetical_statement import HypotheticalStatement

        if not isinstance(hypothetical, HypotheticalStatement):
            raise TypeError("hypothetical must be a HypotheticalStatement")

        causes = [
            self._lemmatize_criteria(getattr(cause, "criteria", {}))
            for cause in (hypothetical.causes or [])
        ]
        effect = self._lemmatize_criteria(getattr(hypothetical.effect, "criteria", {}))

        return HypotheticalStatementExpression(
            causes=causes,
            effect=effect,
            probability=hypothetical.probability,
            validity=hypothetical.validity,
            effect_belief_id=hypothetical.effect_belief_id,
        )

