from typing import List, Optional, Dict, Set, Any
from datetime import datetime, timedelta
import json
from analogai.foundation.modules.belief.belief_repository import BeliefRepository
from analogai.foundation.modules.notion.notion import Notion
from analogai.foundation.modules.belief.belief import Belief
from analogai.foundation.modules.user.user import User
from analogai.foundation.modules.agent.agent import Agent
from analogai.foundation.core.value_objects.modifier import Modifier
from analogai.foundation.common.utils.modifier_utils import ModifierMatcher
from analogai.foundation.common.enums.deontic_modality import DeonticModality
from analogai.foundation.core.value_objects.time import Time
from analogai.foundation.infrastructure.storage.neo4j.neo4j_gateway import Neo4jGateway
from analogai.foundation.infrastructure.mappers import BeliefStorageMapper, NotionStorageMapper, TimeStorageMapper
from analogai.foundation.common.utils.time_resolution import resolve_modifier_to_absolute
import hashlib
from neo4j import Transaction


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 Neo4jBeliefRepositoryImpl(BeliefRepository):
    def __init__(self, belief_service=None):
        self.neo4j_gateway = Neo4jGateway()
        self._belief_cache: Dict[str, Belief] = {}
        self._loading_belief_ids: Set[str] = set()
        self._belief_service = belief_service
        self._matcher = ModifierMatcher()
        self._time_service: Optional[Any] = None

    def set_belief_service(self, belief_service) -> None:
        self._belief_service = belief_service

    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 _generate_unique_id(self, subject_notion: Notion, complement_notion: Notion, relation) -> str:
        from analogai.foundation.common.utils.relationship_type_storage import to_storage_value
        rel_value = to_storage_value(relation)
        unique_string = f"{subject_notion.id}:{complement_notion.id}:{rel_value}"
        return hashlib.md5(unique_string.encode()).hexdigest()

    def _build_belief(self, record: dict) -> Belief:
        belief_data = dict(record["b"])
        subject_data = dict(record["source"])
        object_data = dict(record["target"])
        statement_ids = belief_data.get("statement_ids", [])
        if not statement_ids:
            raise ValueError(f"Cannot build belief {belief_data.get('id', 'unknown')} without statements")

        created_at = datetime.fromisoformat(belief_data.get("created_at", datetime.now().isoformat())) if "created_at" in belief_data else None
        updated_at = datetime.fromisoformat(belief_data.get("updated_at", datetime.now().isoformat())) if "updated_at" in belief_data else None
        user_id = belief_data.get('user_id', belief_data.get('author_id', ''))
        agent_id = belief_data.get('agent_id', '')
        subject_user_id = subject_data.get("user_id", user_id)
        subject_agent_id = subject_data.get("agent_id", agent_id)
        object_user_id = object_data.get("user_id", user_id)
        object_agent_id = object_data.get("agent_id", agent_id)

        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")
        )

        belief_dict = {
            "id": belief_data["id"],
            "agent_id": agent_id,
            "subject_notion": {
                "id": subject_data["id"],
                "user_id": subject_user_id,
                "agent_id": subject_agent_id,
                "caption": subject_data["caption"],
                "context": subject_data.get("context", ""),
                "metadata": subject_data.get("metadata", ""),
                "attributes": subject_attributes,
            },
            "complement_notion": {
                "id": object_data["id"],
                "user_id": object_user_id,
                "agent_id": object_agent_id,
                "caption": object_data["caption"],
                "context": object_data.get("context", ""),
                "metadata": object_data.get("metadata", ""),
                "attributes": object_attributes,
            },
            "relation": belief_data.get("relation") or belief_data.get("relationship_type"),
            "created_at": created_at.isoformat() if created_at else None,
            "updated_at": updated_at.isoformat() if updated_at else None
        }
        if belief_data.get("location"):
            belief_dict["location"] = belief_data.get("location")
        if belief_data.get("from_time_year") is not None:
            belief_dict["from_time_year"] = belief_data.get("from_time_year")
        if belief_data.get("from_time_month") is not None:
            belief_dict["from_time_month"] = belief_data.get("from_time_month")
        if belief_data.get("from_time_day") is not None:
            belief_dict["from_time_day"] = belief_data.get("from_time_day")
        if belief_data.get("from_time_hour") is not None:
            belief_dict["from_time_hour"] = belief_data.get("from_time_hour")
        if belief_data.get("from_time_minute") is not None:
            belief_dict["from_time_minute"] = belief_data.get("from_time_minute")
        if belief_data.get("to_time_year") is not None:
            belief_dict["to_time_year"] = belief_data.get("to_time_year")
        if belief_data.get("to_time_month") is not None:
            belief_dict["to_time_month"] = belief_data.get("to_time_month")
        if belief_data.get("to_time_day") is not None:
            belief_dict["to_time_day"] = belief_data.get("to_time_day")
        if belief_data.get("to_time_hour") is not None:
            belief_dict["to_time_hour"] = belief_data.get("to_time_hour")
        if belief_data.get("to_time_minute") is not None:
            belief_dict["to_time_minute"] = belief_data.get("to_time_minute")
        if belief_data.get("deontic_modality"):
            belief_dict["deontic_modality"] = belief_data.get("deontic_modality")
        if belief_data.get("quantity") is not None:
            belief_dict["quantity"] = belief_data.get("quantity")
        if belief_data.get("measures_json"):
            try:
                belief_dict["attributes"] = json.loads(belief_data.get("measures_json"))
            except (TypeError, ValueError):
                pass

        statements = []
        with self.neo4j_gateway.session() as neo4j_session:
            for stmt_id in statement_ids:
                result = neo4j_session.run(
                    """
                    MATCH (source:Notion)-[s {id: $stmt_id}]->(target:Notion)
                    WHERE type(s) IN ['STATEMENT']
                    RETURN s, source, target, type(s) AS rel_type
                    """,
                    stmt_id=stmt_id
                )
                stmt_record = result.single()
                if not stmt_record:
                    continue

                stmt_data = dict(stmt_record["s"])
                stmt_source = dict(stmt_record["source"])
                stmt_target = dict(stmt_record["target"])
                stmt_type = stmt_record["rel_type"]
                stmt_user_id = stmt_data.get("user_id", stmt_data.get("author_id", ""))
                stmt_agent_id = stmt_data.get('agent_id', '')
                stmt_subject_user_id = stmt_source.get("user_id", stmt_user_id)
                stmt_subject_agent_id = stmt_source.get("agent_id", stmt_agent_id)
                stmt_object_user_id = stmt_target.get("user_id", stmt_user_id)
                stmt_object_agent_id = stmt_target.get("agent_id", stmt_agent_id)

                from analogai.foundation.common.utils.modifier_utils import ModifierBuilder
                modifier = None
                from_time = TimeStorageMapper.unflatten_from_dict(stmt_data, "from_time")
                if not from_time and stmt_data.get("from_time"):
                    from_time = TimeStorageMapper.from_dict(stmt_data.get("from_time"))
                to_time = TimeStorageMapper.unflatten_from_dict(stmt_data, "to_time")
                if not to_time and stmt_data.get("to_time"):
                    to_time = TimeStorageMapper.from_dict(stmt_data.get("to_time"))
                location = stmt_data.get("location")
                deontic_modality = stmt_data.get("deontic_modality")
                
                if any([location, from_time, to_time, deontic_modality]):
                    builder = ModifierBuilder()
                    if location:
                        builder.with_location(location)
                    if from_time:
                        builder.with_from_time(
                            year=from_time.year,
                            month=from_time.month,
                            day=from_time.day,
                            hour=from_time.hour,
                            minute=from_time.minute
                        )
                    if to_time:
                        builder.with_to_time(
                            year=to_time.year,
                            month=to_time.month,
                            day=to_time.day,
                            hour=to_time.hour,
                            minute=to_time.minute
                        )
                    if deontic_modality:
                        builder.with_deontic_modality(DeonticModality(deontic_modality))
                    modifier = builder.build()
                
                stmt_subject_attributes = self._deserialize_attributes(
                    stmt_source.get("attributes_json") or stmt_source.get("attributes")
                )
                stmt_object_attributes = self._deserialize_attributes(
                    stmt_target.get("attributes_json") or stmt_target.get("attributes")
                )

                base_kwargs = {
                    "id": stmt_data["id"],
                    "subject_notion": Notion(
                        user_id=stmt_subject_user_id,
                        agent_id=stmt_subject_agent_id,
                        caption=stmt_source["caption"],
                        id=stmt_source["id"],
                        attributes=stmt_subject_attributes,
                    ),
                    "complement_notion": Notion(
                        user_id=stmt_object_user_id,
                        agent_id=stmt_object_agent_id,
                        caption=stmt_target["caption"],
                        id=stmt_target["id"],
                        attributes=stmt_object_attributes,
                    ),
                    "relation": str(stmt_data.get("relation") or stmt_data.get("relationship_type") or "is"),
                    "modifier": modifier
                }
                
                validity_raw = stmt_data.get("validity")
                statement_kwargs = {
                    **base_kwargs,
                    "probability": stmt_data.get("probability"),
                    "validity": float(validity_raw) if validity_raw is not None else 1.0
                }

                statement = None
                if stmt_type == "STATEMENT":
                    from analogai.foundation.infrastructure.mappers import StatementStorageMapper
                    statement_data = {
                        "id": stmt_data["id"],
                        "user_id": stmt_user_id,
                        "agent_id": stmt_agent_id,
                        "subject_notion": {
                            "id": stmt_source["id"],
                            "user_id": stmt_subject_user_id,
                            "agent_id": stmt_subject_agent_id,
                            "caption": stmt_source["caption"],
                            "attributes": stmt_subject_attributes,
                        },
                        "relation": str(stmt_data.get("relation") or stmt_data.get("relationship_type") or "is"),
                        "complement_notion": {
                            "id": stmt_target["id"],
                            "user_id": stmt_object_user_id,
                            "agent_id": stmt_object_agent_id,
                            "caption": stmt_target["caption"],
                            "attributes": stmt_object_attributes,
                        },
                        "probability": stmt_data.get("probability", 1.0),
                        "validity": stmt_data.get("validity", 1.0),
                    }
                    if stmt_data.get("measures_json"):
                        try:
                            statement_data["attributes"] = json.loads(stmt_data.get("measures_json"))
                        except (TypeError, ValueError):
                            pass
                    if stmt_data.get("quantity") is not None:
                        statement_data["quantity"] = stmt_data["quantity"]
                    statement = StatementStorageMapper.from_dict(statement_data)
                    if stmt_data.get("from_time_month") is not None:
                        statement_data["from_time_month"] = stmt_data.get("from_time_month")
                    if stmt_data.get("from_time_day") is not None:
                        statement_data["from_time_day"] = stmt_data.get("from_time_day")
                    if stmt_data.get("from_time_hour") is not None:
                        statement_data["from_time_hour"] = stmt_data.get("from_time_hour")
                    if stmt_data.get("from_time_minute") is not None:
                        statement_data["from_time_minute"] = stmt_data.get("from_time_minute")
                    if stmt_data.get("to_time_year") is not None:
                        statement_data["to_time_year"] = stmt_data.get("to_time_year")
                    if stmt_data.get("to_time_month") is not None:
                        statement_data["to_time_month"] = stmt_data.get("to_time_month")
                    if stmt_data.get("to_time_day") is not None:
                        statement_data["to_time_day"] = stmt_data.get("to_time_day")
                    if stmt_data.get("to_time_hour") is not None:
                        statement_data["to_time_hour"] = stmt_data.get("to_time_hour")
                    if stmt_data.get("to_time_minute") is not None:
                        statement_data["to_time_minute"] = stmt_data.get("to_time_minute")
                    if stmt_data.get("deontic_modality"):
                        statement_data["deontic_modality"] = stmt_data["deontic_modality"]
                    statement = StatementStorageMapper.from_dict(statement_data)
                else:
                    raise ValueError(f"Invalid statement type {stmt_type}")

                if statement:
                    statements.append(statement)

        if not statements:
            raise ValueError(f"Cannot build belief {belief_data['id']} - no valid statements loaded")

        belief = BeliefStorageMapper.from_dict(belief_dict, statements=statements)
        return belief

    def delete(self, belief: Belief) -> None:
        def delete_belief(tx: Transaction, belief: Belief) -> int:
            result = tx.run(
                """
                MATCH (source:Notion {id: $subject_id})-[b:BELIEF {id: $id}]->(target:Notion {id: $object_id})
                DELETE b
                RETURN count(b) as deleted_count
                """,
                subject_id=belief.subject_notion.id,
                id=belief.id,
                object_id=belief.complement_notion.id
            )
            return result.single()["deleted_count"]

        with self.neo4j_gateway.session() as neo4j_session:
            deleted_count = neo4j_session.execute_write(delete_belief, belief)
            if deleted_count == 0:
                raise ValueError(f"No belief with id {belief.id} was deleted")

    def _belief_has_matching_statement(self, belief: Belief, modifier: Modifier) -> bool:
        return any(self._matcher.matches(statement, modifier) for statement in belief.statements)

    def _search_modifier(self, modifier: Optional[Modifier]) -> Optional[Modifier]:
        if modifier is None:
            return None
        return Modifier(
            location=modifier.location,
            from_time=modifier.from_time,
            to_time=modifier.to_time,
            attributes=modifier.attributes,
            quantity=modifier.quantity,
        )

    def _belief_has_no_modal_statement(self, belief: Belief) -> bool:
        return any(stmt.modifier is None for stmt in belief.statements)

    def find(self, subject_notion: Notion, complement_notion: Notion, relation,
             modifier: Optional[Modifier] = None, quantity: Optional[float] = None) -> Optional[Belief]:
        if quantity is not None:
            if modifier is None:
                modifier = Modifier(quantity=quantity)
            else:
                modifier = Modifier(
                    location=modifier.location,
                    deontic_modality=modifier.deontic_modality,
                    from_time=modifier.from_time,
                    to_time=modifier.to_time,
                    attributes=list(modifier.attributes) if modifier.attributes else None,
                    quantity=quantity,
                )
        from analogai.foundation.common.utils.relationship_type_storage import to_storage_value
        rel_value = to_storage_value(relation)
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion {id: $subject_id})-[b:BELIEF]->(target:Notion {id: $object_id})
                WHERE b.relation = $rel_value OR b.relation = $rel_value
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """,
                subject_id=subject_notion.id,
                object_id=complement_notion.id,
                rel_value=rel_value
            )
            for record in result:
                belief = self._build_belief(record)
                if not belief.statements:
                    raise ValueError(f"Loaded belief {belief.id} has no statements")
                if modifier is None:
                    return belief
                if self._belief_has_matching_statement(belief, modifier):
                    return belief
            return None

    def find_by_id(self, id: str) -> Optional[Belief]:
        if id in self._belief_cache:
            return self._belief_cache[id]
        
        if id in self._loading_belief_ids:
            return None
        
        self._loading_belief_ids.add(id)
        try:
            with self.neo4j_gateway.session() as neo4j_session:
                result = neo4j_session.run(
                    """
                    MATCH (source:Notion)-[b:BELIEF {id: $id}]->(target:Notion)
                    RETURN b, source, target
                    """,
                    id=id
                )
                record = result.single()
                if not record:
                    return None
                belief = self._build_belief(record)
                if not belief.statements:
                    raise ValueError(f"Loaded belief {id} has no statements")
                self._belief_cache[id] = belief
                return belief
        finally:
            self._loading_belief_ids.discard(id)

    def save(self, belief: Belief) -> Belief:
        if not belief.statements:
            raise ValueError(f"Cannot save belief {belief.id} without statements")
        data = BeliefStorageMapper.to_dict(belief)
        if self._time_service and (belief.from_time or belief.to_time):
            _strip_time_keys(data)
            modifier = Modifier(
                location=belief.location,
                deontic_modality=belief.deontic_modality,
                from_time=belief.from_time,
                to_time=belief.to_time,
            )
            resolved = resolve_modifier_to_absolute(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"))
        with self.neo4j_gateway.session() as neo4j_session:
            set_clauses = [
                "b.statement_ids = $statement_ids",
                "b.created_at = $created_at",
                "b.updated_at = $updated_at"
            ]
            params = {
                "id": data["id"],
                "relation": data["relation"],
                "subject_id": data["subject_notion"]["id"],
                "object_id": data["complement_notion"]["id"],
                "statement_ids": data["statement_ids"],
                "agent_id": data["agent_id"],
                "created_at": data["created_at"],
                "updated_at": data["updated_at"],
            }
            
            if data.get("location"):
                set_clauses.append("b.location = $location")
                params["location"] = data["location"]
            if data.get("from_time_year") is not None:
                set_clauses.append("b.from_time_year = $from_time_year")
                params["from_time_year"] = data.get("from_time_year")
            if data.get("from_time_month") is not None:
                set_clauses.append("b.from_time_month = $from_time_month")
                params["from_time_month"] = data.get("from_time_month")
            if data.get("from_time_day") is not None:
                set_clauses.append("b.from_time_day = $from_time_day")
                params["from_time_day"] = data.get("from_time_day")
            if data.get("from_time_hour") is not None:
                set_clauses.append("b.from_time_hour = $from_time_hour")
                params["from_time_hour"] = data.get("from_time_hour")
            if data.get("from_time_minute") is not None:
                set_clauses.append("b.from_time_minute = $from_time_minute")
                params["from_time_minute"] = data.get("from_time_minute")
            if data.get("to_time_year") is not None:
                set_clauses.append("b.to_time_year = $to_time_year")
                params["to_time_year"] = data.get("to_time_year")
            if data.get("to_time_month") is not None:
                set_clauses.append("b.to_time_month = $to_time_month")
                params["to_time_month"] = data.get("to_time_month")
            if data.get("to_time_day") is not None:
                set_clauses.append("b.to_time_day = $to_time_day")
                params["to_time_day"] = data.get("to_time_day")
            if data.get("to_time_hour") is not None:
                set_clauses.append("b.to_time_hour = $to_time_hour")
                params["to_time_hour"] = data.get("to_time_hour")
            if data.get("to_time_minute") is not None:
                set_clauses.append("b.to_time_minute = $to_time_minute")
                params["to_time_minute"] = data.get("to_time_minute")
            if data.get("deontic_modality"):
                set_clauses.append("b.deontic_modality = $deontic_modality")
                params["deontic_modality"] = data["deontic_modality"]
            if data.get("quantity") is not None:
                set_clauses.append("b.quantity = $quantity")
                params["quantity"] = data["quantity"]
            else:
                set_clauses.append("b.quantity = null")
            if data.get("attributes") is not None:
                set_clauses.append("b.measures_json = $measures_json")
                params["measures_json"] = json.dumps(data.get("attributes"))
            else:
                set_clauses.append("b.measures_json = null")
            
            result = neo4j_session.run(
                f"""
                MATCH (source:Notion {{id: $subject_id}})
                MATCH (target:Notion {{id: $object_id}})
                MERGE (source)-[b:BELIEF {{id: $id, relation: $relation, agent_id: $agent_id}}]->(target)
                SET {', '.join(set_clauses)}
                """,
                **params
            )
            result.consume()
        self.clear_cache(belief.id)
        self._belief_cache[belief.id] = belief
        return belief

    def clear_cache(self, belief_id: Optional[str] = None) -> None:
        if belief_id:
            self._belief_cache.pop(belief_id, None)
        else:
            self._belief_cache.clear()

    def search_by_notions(self, subject_notion: Notion, complement_notion: Notion,
                          modifier: Optional[Modifier] = None) -> List[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion {id: $subject_id})-[b:BELIEF]->(target:Notion {id: $object_id})
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """,
                subject_id=subject_notion.id,
                object_id=complement_notion.id
            )
            search_mod = self._search_modifier(modifier)
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief:
                    if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                        beliefs.append(belief)
            return beliefs

    def search_by_subject_notion(self, subject_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion {id: $subject_id})-[b:BELIEF]->(target:Notion)
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """,
                subject_id=subject_notion.id
            )
            search_mod = self._search_modifier(modifier)
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief:
                    if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                        beliefs.append(belief)
            return beliefs

    def search_by_complement_notion(self, complement_notion: Notion,
                                modifier: Optional[Modifier] = None) -> List[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion)-[b:BELIEF]->(target:Notion {id: $object_id})
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """,
                object_id=complement_notion.id
            )
            search_mod = self._search_modifier(modifier)
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief:
                    if not search_mod or self._belief_has_matching_statement(belief, search_mod):
                        beliefs.append(belief)
            return beliefs

    def find_by_agent_id(self, agent_id: str) -> List[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion)-[b:BELIEF {agent_id: $agent_id}]->(target:Notion)
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """,
                agent_id=agent_id
            )
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief:
                    beliefs.append(belief)
            return beliefs

    def get_recently_updated(self, agent_id: str, since: datetime) -> List[Belief]:
        since_iso = since.isoformat()
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion)-[b:BELIEF {agent_id: $agent_id}]->(target:Notion)
                WHERE (b.updated_at >= $since OR b.created_at >= $since)
                RETURN b, source, target
                """,
                agent_id=agent_id,
                since=since_iso
            )
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief:
                    beliefs.append(belief)
            return beliefs

    def search_by_modifier(self, modifier: Modifier) -> List[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion)-[b:BELIEF]->(target:Notion)
                RETURN b, source, target
                ORDER BY b.updated_at DESC, b.id DESC
                """
            )
            beliefs = []
            for record in result:
                belief = self._build_belief(record)
                if belief and self._belief_has_matching_statement(belief, modifier):
                    beliefs.append(belief)
            return beliefs

    def find_belief_by_statement_id(self, statement_id: str) -> Optional[Belief]:
        with self.neo4j_gateway.session() as neo4j_session:
            result = neo4j_session.run(
                """
                MATCH (source:Notion)-[b:BELIEF]->(target:Notion)
                WHERE $statement_id IN b.statement_ids
                RETURN b, source, target
                LIMIT 1
                """,
                statement_id=statement_id
            )
            record = result.single()
            if not record:
                return None
            return self._build_belief(record)