from datetime import datetime
from typing import List, Optional

from analogai.foundation.modules.conditionals.hypothetical_statement import HypotheticalStatement
from analogai.foundation.modules.conditionals.hypothetical_statement_repository import HypotheticalStatementRepository
from analogai.foundation.infrastructure.storage.mongodb.mongodb_storage_adapter_impl import MongoDBStorageAdapterImpl
from analogai.foundation.core.value_objects.cause import Cause
from analogai.foundation.core.value_objects.effect import Effect
from analogai.foundation.common.enums.criterion import Criterion
from analogai.foundation.core.value_objects.relation import Relation


class MongoDBHypotheticalStatementRepositoryImpl(HypotheticalStatementRepository):
    COLLECTION = "hypothetical_statements"

    def __init__(self, storage_adapter: Optional[MongoDBStorageAdapterImpl] = None):
        self.storage_adapter = storage_adapter if storage_adapter else MongoDBStorageAdapterImpl()

    def create(self, hypothetical_statement: HypotheticalStatement) -> HypotheticalStatement:
        data = self._to_dict(hypothetical_statement)
        self.storage_adapter.store(self.COLLECTION, data)
        return hypothetical_statement

    def find_by_id(self, hypothetical_statement_id: str) -> Optional[HypotheticalStatement]:
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION, hypothetical_statement_id)
        if not data:
            return None
        return self._from_dict(data)

    def find_by_agent_id(self, agent_id: str) -> List[HypotheticalStatement]:
        results = self.storage_adapter.retrieve_by_attributes(
            self.COLLECTION,
            {"agent_id": agent_id}
        )
        return [self._from_dict(data) for data in results]

    def _to_dict(self, hypothetical_statement: HypotheticalStatement) -> dict:
        causes_serialized = []
        for cause in hypothetical_statement.causes:
            clause = {}
            for criterion, value in cause.criteria.items():
                try:
                    key = criterion.value
                except AttributeError:
                    continue
                try:
                    clause[key] = value.value
                except AttributeError:
                    clause[key] = value
            causes_serialized.append(clause)

        effect_serialized = {}
        for criterion, value in hypothetical_statement.effect.criteria.items():
            try:
                key = criterion.value
            except AttributeError:
                continue
            try:
                effect_serialized[key] = value.value
            except AttributeError:
                effect_serialized[key] = value

        return {
            "id": hypothetical_statement.id,
            "user_id": hypothetical_statement.user_id,
            "agent_id": hypothetical_statement.agent_id,
            "validity": hypothetical_statement.validity,
            "created_at": hypothetical_statement.created_at.isoformat()
            if hypothetical_statement.created_at
            else None,
            "causes": causes_serialized,
            "effect": effect_serialized,
        }

    def _from_dict(self, data: dict) -> HypotheticalStatement:
        causes_raw = data.get("causes", [])
        causes = [Cause(self._deserialize_criteria(raw or {})) for raw in causes_raw]
        return HypotheticalStatement(
            id=data.get("id"),
            user_id=data.get("user_id"),
            agent_id=data["agent_id"],
            causes=causes,
            effect=Effect(self._deserialize_criteria(data.get("effect") or {})),
            validity=data.get("validity", 1.0),
            created_at=self._deserialize_created_at(data.get("created_at")),
        )

    def _deserialize_created_at(self, value) -> Optional[datetime]:
        if value is None:
            return None
        if isinstance(value, datetime):
            return value
        try:
            return datetime.fromisoformat(value)
        except (TypeError, ValueError):
            return None

    def _deserialize_criteria(self, raw: dict) -> dict:
        criteria = {}
        for key, value in raw.items():
            try:
                criterion = Criterion(key)
            except ValueError:
                continue
            if criterion == Criterion.Relation:
                criteria[criterion] = Relation.from_string(value) if value else None
            else:
                if type(value) is list:
                    criteria[criterion] = tuple(value)
                else:
                    criteria[criterion] = value
        return criteria
