from datetime import datetime
from typing import List, Optional, Any
import json

from analogai.foundation.modules.direct_statement.direct_statement import DirectStatement
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.direct_statement.direct_statement_repository import StatementRepository
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.core.value_objects.relation import Relation
from analogai.foundation.infrastructure.mappers.statement_storage_mapper import StatementStorageMapper
from analogai.foundation.infrastructure.mappers.time_storage_mapper import TimeStorageMapper
from analogai.foundation.common.utils.modifier_utils import ModifierMatcher
from analogai.foundation.common.utils.time_resolution import resolve_modifier_to_absolute
from analogai.foundation.infrastructure.storage.neo4j.neo4j_gateway import Neo4jGateway


def _strip_time_keys(data: dict) -> None:
    time_keys = [key for key in list(data) if key.startswith("from_time_") or key.startswith("to_time_")]
    for key in time_keys:
        del data[key]


class Neo4jStatementRepositoryImpl(StatementRepository):
    def __init__(self):
        self.neo4j_gateway = Neo4jGateway()
        self._modal_matcher = ModifierMatcher()
        self._time_service: Optional[Any] = None

    def set_time_service(self, time_service: Any) -> None:
        self._time_service = time_service

    def _deserialize_attributes(self, raw: Optional[object]) -> Optional[list]:
        if raw is None:
            return None
        if isinstance(raw, list):
            return raw
        if isinstance(raw, str):
            try:
                return json.loads(raw)
            except (TypeError, ValueError):
                return None
        return None

    def _modal_key_from_data(self, data: dict) -> str:
        loc = data.get("location")
        deontic = data.get("deontic_modality")
        fy = data.get("from_time_year")
        fm = data.get("from_time_month")
        fd = data.get("from_time_day")
        fh = data.get("from_time_hour")
        fmi = data.get("from_time_minute")
        ty = data.get("to_time_year")
        tmo = data.get("to_time_month")
        td = data.get("to_time_day")
        th = data.get("to_time_hour")
        tmi = data.get("to_time_minute")
        if loc is None and deontic is None and fy is None and fm is None and fd is None and fh is None and fmi is None and ty is None and tmo is None and td is None and th is None and tmi is None:
            return "__none__"
        from_part = "_".join(str(x) if x is not None else "" for x in (fy, fm, fd, fh, fmi))
        to_part = "_".join(str(x) if x is not None else "" for x in (ty, tmo, td, th, tmi))
        return f"loc={loc or ''}|deontic={deontic or ''}|from={from_part}|to={to_part}"

    def _modal_key_from_modifier(self, modifier: Optional[Modifier]) -> str:
        if modifier is None:
            return "__none__"
        data = {"location": modifier.location, "deontic_modality": modifier.deontic_modality.value if modifier.deontic_modality else None}
        if modifier.from_time:
            data.update(TimeStorageMapper.flatten_to_dict(modifier.from_time, "from_time"))
        if modifier.to_time:
            data.update(TimeStorageMapper.flatten_to_dict(modifier.to_time, "to_time"))
        return self._modal_key_from_data(data)

    def create(self, statement: DirectStatement) -> DirectStatement:
        data = StatementStorageMapper.to_dict(statement)
        if self._time_service and statement.modifier:
            _strip_time_keys(data)
            resolved = resolve_modifier_to_absolute(statement.modifier, self._time_service)
            if resolved:
                data.update(TimeStorageMapper.flatten_to_dict_absolute_only(resolved.from_time, "from_time"))
                data.update(TimeStorageMapper.flatten_to_dict_absolute_only(resolved.to_time, "to_time"))
        modal_key = self._modal_key_from_data(data)
        created_at = data.get("created_at")
        if created_at is not None:
            try:
                created_at = datetime.fromisoformat(created_at)
            except TypeError:
                pass
        if created_at is None:
            created_at = datetime.now()
        relation = data.get("relation", "")
        params = {
            "id": data["id"],
            "subject_id": data["subject_notion"]["id"],
            "object_id": data["complement_notion"]["id"],
            "user_id": data["user_id"],
            "agent_id": data["agent_id"],
            "relation": relation,
            "modal_key": modal_key,
            "created_at": created_at.isoformat(),
            "probability": data.get("probability", 1.0),
            "validity": data.get("validity", 1.0),
        }
        if data.get("quantity") is not None:
            params["quantity"] = data["quantity"]
        if data.get("location"):
            params["location"] = data["location"]
        if data.get("from_time_year") is not None:
            params["from_time_year"] = data.get("from_time_year")
        if data.get("from_time_month") is not None:
            params["from_time_month"] = data.get("from_time_month")
        if data.get("from_time_day") is not None:
            params["from_time_day"] = data.get("from_time_day")
        if data.get("from_time_hour") is not None:
            params["from_time_hour"] = data.get("from_time_hour")
        if data.get("from_time_minute") is not None:
            params["from_time_minute"] = data.get("from_time_minute")
        if data.get("to_time_year") is not None:
            params["to_time_year"] = data.get("to_time_year")
        if data.get("to_time_month") is not None:
            params["to_time_month"] = data.get("to_time_month")
        if data.get("to_time_day") is not None:
            params["to_time_day"] = data.get("to_time_day")
        if data.get("to_time_hour") is not None:
            params["to_time_hour"] = data.get("to_time_hour")
        if data.get("to_time_minute") is not None:
            params["to_time_minute"] = data.get("to_time_minute")
        if data.get("deontic_modality"):
            params["deontic_modality"] = data["deontic_modality"]
        if data.get("attributes") is not None:
            params["measures_json"] = json.dumps(data.get("attributes"))

        edge_props = ["id: $id", "user_id: $user_id", "agent_id: $agent_id", "relation: $relation", "modal_key: $modal_key", "created_at: $created_at", "probability: $probability", "validity: $validity"]
        if data.get("quantity") is not None:
            edge_props.append("quantity: $quantity")
        if data.get("location"):
            edge_props.append("location: $location")
        if data.get("from_time_year") is not None:
            edge_props.append("from_time_year: $from_time_year")
        if data.get("from_time_month") is not None:
            edge_props.append("from_time_month: $from_time_month")
        if data.get("from_time_day") is not None:
            edge_props.append("from_time_day: $from_time_day")
        if data.get("from_time_hour") is not None:
            edge_props.append("from_time_hour: $from_time_hour")
        if data.get("from_time_minute") is not None:
            edge_props.append("from_time_minute: $from_time_minute")
        if data.get("to_time_year") is not None:
            edge_props.append("to_time_year: $to_time_year")
        if data.get("to_time_month") is not None:
            edge_props.append("to_time_month: $to_time_month")
        if data.get("to_time_day") is not None:
            edge_props.append("to_time_day: $to_time_day")
        if data.get("to_time_hour") is not None:
            edge_props.append("to_time_hour: $to_time_hour")
        if data.get("to_time_minute") is not None:
            edge_props.append("to_time_minute: $to_time_minute")
        if data.get("deontic_modality"):
            edge_props.append("deontic_modality: $deontic_modality")
        if data.get("attributes") is not None:
            edge_props.append("measures_json: $measures_json")

        with self.neo4j_gateway.session() as session:
            session.run(
                f"""
                MERGE (subject:Notion {{id: $subject_id}})
                MERGE (object:Notion {{id: $object_id}})
                CREATE (subject)-[s:STATEMENT {{{", ".join(edge_props)}}}]->(object)
                """,
                **params
            )
        return statement

    def update(self, statement: DirectStatement) -> DirectStatement:
        data = StatementStorageMapper.to_dict(statement)
        if self._time_service and statement.modifier:
            _strip_time_keys(data)
            resolved = resolve_modifier_to_absolute(statement.modifier, self._time_service)
            if resolved:
                data.update(TimeStorageMapper.flatten_to_dict_absolute_only(resolved.from_time, "from_time"))
                data.update(TimeStorageMapper.flatten_to_dict_absolute_only(resolved.to_time, "to_time"))
        modal_key = self._modal_key_from_data(data)
        params = {
            "id": data["id"],
            "subject_id": data["subject_notion"]["id"],
            "object_id": data["complement_notion"]["id"],
            "probability": data.get("probability", 1.0),
            "validity": data.get("validity", 1.0),
            "modal_key": modal_key,
        }
        set_parts = ["s.probability = $probability", "s.validity = $validity", "s.modal_key = $modal_key"]
        if data.get("quantity") is not None:
            params["quantity"] = data["quantity"]
            set_parts.append("s.quantity = $quantity")
        else:
            set_parts.append("s.quantity = null")
        if data.get("location"):
            params["location"] = data["location"]
            set_parts.append("s.location = $location")
        else:
            set_parts.append("s.location = null")
        for key in ["from_time_year", "from_time_month", "from_time_day", "from_time_hour", "from_time_minute"]:
            if data.get(key) is not None:
                params[key] = data[key]
                set_parts.append(f"s.{key} = ${key}")
            else:
                set_parts.append(f"s.{key} = null")
        for key in ["to_time_year", "to_time_month", "to_time_day", "to_time_hour", "to_time_minute"]:
            if data.get(key) is not None:
                params[key] = data[key]
                set_parts.append(f"s.{key} = ${key}")
            else:
                set_parts.append(f"s.{key} = null")
        if data.get("deontic_modality"):
            params["deontic_modality"] = data["deontic_modality"]
            set_parts.append("s.deontic_modality = $deontic_modality")
        else:
            set_parts.append("s.deontic_modality = null")
        if data.get("attributes") is not None:
            params["measures_json"] = json.dumps(data.get("attributes"))
            set_parts.append("s.measures_json = $measures_json")
        else:
            set_parts.append("s.measures_json = null")
        with self.neo4j_gateway.session() as session:
            session.run(
                f"""
                MATCH (subject:Notion {{id: $subject_id}})-[s:STATEMENT {{id: $id}}]->(object:Notion {{id: $object_id}})
                SET {", ".join(set_parts)}
                """,
                **params
            )
        return statement

    def _record_to_dict(self, record) -> dict:
        edge = dict(record["s"])
        subject_data = dict(record["subject"])
        object_data = dict(record["object"])
        subject_attributes = self._deserialize_attributes(
            subject_data.get("attributes_json") or subject_data.get("attributes")
        )
        object_attributes = self._deserialize_attributes(
            object_data.get("attributes_json") or object_data.get("attributes")
        )
        relation = edge.get("relation") or edge.get("relationship", "")
        created_at = None
        if edge.get("created_at"):
            raw = edge["created_at"]
            try:
                created_at = datetime.fromisoformat(raw)
            except TypeError:
                created_at = raw
        result = {
            "id": edge["id"],
            "user_id": edge["user_id"],
            "agent_id": edge["agent_id"],
            "created_at": created_at,
            "probability": edge.get("probability", 1.0),
            "validity": edge.get("validity", 1.0),
            "subject_notion": {
                "id": subject_data["id"],
                "user_id": subject_data.get("user_id", ""),
                "agent_id": subject_data.get("agent_id", ""),
                "caption": subject_data.get("caption", ""),
                "attributes": subject_attributes,
            },
            "relation": relation,
            "complement_notion": {
                "id": object_data["id"],
                "user_id": object_data.get("user_id", ""),
                "agent_id": object_data.get("agent_id", ""),
                "caption": object_data.get("caption", ""),
                "attributes": object_attributes,
            },
        }
        if edge.get("quantity") is not None:
            result["quantity"] = edge["quantity"]
        if edge.get("location"):
            result["location"] = edge["location"]
        if edge.get("from_time_year") is not None:
            result["from_time_year"] = edge.get("from_time_year")
        if edge.get("from_time_month") is not None:
            result["from_time_month"] = edge.get("from_time_month")
        if edge.get("from_time_day") is not None:
            result["from_time_day"] = edge.get("from_time_day")
        if edge.get("from_time_hour") is not None:
            result["from_time_hour"] = edge.get("from_time_hour")
        if edge.get("from_time_minute") is not None:
            result["from_time_minute"] = edge.get("from_time_minute")
        if edge.get("to_time_year") is not None:
            result["to_time_year"] = edge.get("to_time_year")
        if edge.get("to_time_month") is not None:
            result["to_time_month"] = edge.get("to_time_month")
        if edge.get("to_time_day") is not None:
            result["to_time_day"] = edge.get("to_time_day")
        if edge.get("to_time_hour") is not None:
            result["to_time_hour"] = edge.get("to_time_hour")
        if edge.get("to_time_minute") is not None:
            result["to_time_minute"] = edge.get("to_time_minute")
        if edge.get("deontic_modality"):
            result["deontic_modality"] = edge["deontic_modality"]
        raw_measures = edge.get("measures_json")
        if raw_measures:
            try:
                result["attributes"] = json.loads(raw_measures)
            except (TypeError, ValueError):
                pass
        return result

    def find_by_id(self, statement_id: str) -> Optional[DirectStatement]:
        with self.neo4j_gateway.session() as session:
            result = session.run(
                """
                MATCH (subject:Notion)-[s:STATEMENT {id: $id}]->(object:Notion)
                RETURN s, subject, object
                """,
                id=statement_id,
            )
            record = result.single()
            if not record:
                return None
            data = self._record_to_dict(record)
            return StatementStorageMapper.from_dict(data)

    def find_by_subject(self, subject_notion: Notion) -> List[DirectStatement]:
        with self.neo4j_gateway.session() as session:
            result = session.run(
                """
                MATCH (subject:Notion {id: $subject_id})-[s:STATEMENT]->(object:Notion)
                RETURN s, subject, object
                """,
                subject_id=subject_notion.id,
            )
            statements = []
            for rec in result:
                data = self._record_to_dict(rec)
                statements.append(StatementStorageMapper.from_dict(data))
            return statements

    def find_by_subject_and_relation(
        self,
        subject_notion: Notion,
        relation: Relation,
        modifier: Optional[Modifier] = None,
    ) -> List[DirectStatement]:
        statements = self.find_by_subject(subject_notion)
        by_type = [p for p in statements if p.relation == relation]
        if modifier is None:
            return by_type
        return [p for p in by_type if self._modal_matcher.context_matches(p.modifier, modifier)]

    def find_by_object(self, complement_notion: Notion) -> List[DirectStatement]:
        with self.neo4j_gateway.session() as session:
            result = session.run(
                """
                MATCH (subject:Notion)-[s:STATEMENT]->(object:Notion {id: $object_id})
                RETURN s, subject, object
                """,
                object_id=complement_notion.id,
            )
            statements = []
            for rec in result:
                data = self._record_to_dict(rec)
                statements.append(StatementStorageMapper.from_dict(data))
            return statements

    def find_by_object_and_relation(
        self,
        complement_notion: Notion,
        relation: Relation,
        modifier: Optional[Modifier] = None,
    ) -> List[DirectStatement]:
        statements = self.find_by_object(complement_notion)
        by_type = [p for p in statements if p.relation == relation]
        if modifier is None:
            return by_type
        return [p for p in by_type if self._modal_matcher.context_matches(p.modifier, modifier)]

    def find_by_agent_id(
        self,
        agent_id: str,
        limit: int = 0,
        offset: int = 0,
    ) -> List[DirectStatement]:
        query = """
                MATCH (subject:Notion)-[s:STATEMENT {agent_id: $agent_id}]->(object:Notion)
                RETURN s, subject, object
                SKIP $offset
                """
        params = {
            "agent_id": agent_id,
            "offset": offset,
        }
        if limit > 0:
            query += "\n                LIMIT $limit"
            params["limit"] = limit
        with self.neo4j_gateway.session() as session:
            result = session.run(
                query,
                **params,
            )
            statements = []
            for rec in result:
                data = self._record_to_dict(rec)
                statements.append(StatementStorageMapper.from_dict(data))
            return statements

    def find_direct_statement(
        self,
        user_id: str,
        agent_id: str,
        subject_notion: Notion,
        relation: Relation,
        complement_notion: Notion,
        modifier: Optional[Modifier] = None,
    ) -> Optional[DirectStatement]:
        statements = self.find_by_subject_and_relation(
            subject_notion, relation, modifier
        )
        for stmt in statements:
            if (
                stmt.complement_notion.id == complement_notion.id
                and stmt.user_id == user_id
                and stmt.agent_id == agent_id
            ):
                return stmt
        return None