import json
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.neo4j.neo4j_gateway import Neo4jGateway


class Neo4jHypotheticalStatementRepositoryImpl(HypotheticalStatementRepository):
    LABEL = "hypothetical_statement"

    def __init__(self):
        self.neo4j_gateway = Neo4jGateway()

    def create(self, hypothetical_statement: HypotheticalStatement) -> HypotheticalStatement:
        causes_serialized = self._serialize_causes(hypothetical_statement.causes)
        with self.neo4j_gateway.session() as session:
            session.run(
                f"CREATE (c:{self.LABEL} {{id: $id, agent_id: $agent_id, effect_belief_id: $effect_belief_id, "
                "probability: $probability, causes: $causes})",
                id=hypothetical_statement.id,
                agent_id=hypothetical_statement.agent_id,
                effect_belief_id=hypothetical_statement.effect_belief_id,
                probability=hypothetical_statement.probability,
                causes=causes_serialized,
            )
        return hypothetical_statement

    def find_by_id(self, hypothetical_statement_id: str) -> Optional[HypotheticalStatement]:
        with self.neo4j_gateway.session() as session:
            result = session.run(
                f"MATCH (c:{self.LABEL} {{id: $id}}) RETURN c",
                id=hypothetical_statement_id,
            )
            record = result.single()
            if not record:
                return None
            return self._from_node(record["c"])

    def find_by_agent_id(self, agent_id: str) -> List[HypotheticalStatement]:
        with self.neo4j_gateway.session() as session:
            result = session.run(
                f"MATCH (c:{self.LABEL} {{agent_id: $agent_id}}) RETURN c",
                agent_id=agent_id,
            )
            return [self._from_node(record["c"]) for record in result]

    def _serialize_causes(self, causes) -> str:
        serialized = []
        for c in causes:
            clause = {}
            for k, v in c.criteria.items():
                try:
                    key = k.value
                except AttributeError:
                    continue
                try:
                    clause[key] = v.value
                except AttributeError:
                    clause[key] = v
            serialized.append(clause)
        return json.dumps(serialized)

    def _from_node(self, node) -> HypotheticalStatement:
        from analogai.foundation.core.value_objects.cause import Cause

        causes_raw = json.loads(node.get("causes", "[]"))
        causes = [Cause(self._deserialize_criteria(raw or {})) for raw in causes_raw]
        return HypotheticalStatement(
            id=node.get("id"),
            agent_id=node["agent_id"],
            causes=causes,
            effect_belief_id=node.get("effect_belief_id", ""),
            probability=node.get("probability", 1.0),
        )

    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:
                continue
            if criterion == Criterion.Relation:
                criteria[criterion] = str(value) if value is not None else value
            else:
                if type(value) is list:
                    criteria[criterion] = tuple(value)
                else:
                    criteria[criterion] = value
        return criteria
