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.in_memory.in_memory_storage_adapter_impl import InMemoryStorageAdapterImpl
from analogai.foundation.core.value_objects.relation import Relation


class InMemoryHypotheticalStatementRepositoryImpl(HypotheticalStatementRepository):
    COLLECTION_NAME = "hypothetical_statements"

    def __init__(self, storage_adapter=None):
        self.storage_adapter = storage_adapter if storage_adapter else InMemoryStorageAdapterImpl()

    def create(self, hypothetical_statement: HypotheticalStatement) -> HypotheticalStatement:
        def _serialize_clause(clause: dict) -> dict:
            serialized = {}
            for k, v in clause.items():
                try:
                    key = k.value
                except AttributeError:
                    if isinstance(k, str):
                        key = k
                    else:
                        continue
                try:
                    serialized[key] = v.value
                except AttributeError:
                    serialized[key] = v
            return serialized

        causes_serialized = [_serialize_clause(c.criteria) for c in hypothetical_statement.causes]
        effect_serialized = _serialize_clause(hypothetical_statement.effect.criteria)
        data = {
            "id": hypothetical_statement.id,
            "user_id": hypothetical_statement.user_id,
            "agent_id": hypothetical_statement.agent_id,
            "causes": causes_serialized,
            "effect": effect_serialized,
            "validity": hypothetical_statement.validity,
            "created_at": hypothetical_statement.created_at.isoformat()
            if hypothetical_statement.created_at
            else None,
        }
        self.storage_adapter.store(self.COLLECTION_NAME, data)
        return hypothetical_statement

    def find_by_id(self, hypothetical_statement_id: str) -> Optional[HypotheticalStatement]:
        data = self.storage_adapter.retrieve_by_id(self.COLLECTION_NAME, 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_NAME, {"agent_id": agent_id}
        )
        return [self._from_dict(d) for d in results]

    def _from_dict(self, data: dict) -> HypotheticalStatement:
        from analogai.foundation.core.value_objects.cause import Cause
        from analogai.foundation.core.value_objects.effect import Effect
        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:
        from analogai.foundation.common.enums.criterion import Criterion
        criteria = {}
        for key, value in raw.items():
            try:
                criterion = Criterion(key)
            except ValueError:
                criterion = None
            if criterion is None:
                if key == "relation":
                    criteria[key] = Relation.from_string(value) if value else None
                else:
                    criteria[key] = tuple(value) if type(value) is list else value
                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
