"""
neo4j_store.py
---------------
Neo4j-backed symbolic persistence layer for HIEROMEM.

Persists the symbolic scene-graph (G) component of each Episode,
enabling graph-level reasoning, relationship queries, and visualization.

Usage:
    import os
    store = Neo4jStore(
        uri=os.environ["NEO4J_URI"],
        user=os.environ["NEO4J_USER"],
        password=os.environ["NEO4J_PASSWORD"])
    store.save_episode(ep)
    store.query_relations("cat")

Author: Antigravity
"""

from __future__ import annotations
from .neo4j_tenant_manager import Neo4jTenantManager
from typing import Iterable, Optional, Any

from datetime import datetime

import networkx as nx
import numpy as np
from neo4j import GraphDatabase

from hieromem_core.episode import Episode
from hieromem_core.circuit_breaker import CircuitBreaker
from hieromem_core.utils.retry import with_retry
from neo4j.exceptions import ServiceUnavailable, SessionExpired, TransientError
from .graph_store import GraphStore


class Neo4jStore(GraphStore):
    """
    Persistence backend for the symbolic graph layer of HIEROMEM.

    Relation Normalization:
        This store normalizes relation keys to match InMemoryGraphStore:
        - Reads from: ``relation_type``, ``relation``, or ``rel``
        - Neo4j relationship types are UPPERCASE (e.g., ``OWNS``)
        - Stored properties include both ``relation`` and ``relation_type`` (lowercase)
        - This ensures compatibility with GSR and cross-backend queries

    Concurrency & Connection Handling:
        - Uses Neo4j's built-in connection pooling via ``GraphDatabase.driver``
        - Connection pool is managed automatically with keep-alive
        - CircuitBreaker prevents cascading failures (5 failures = 30s cooldown)
        - Retry logic handles transient errors (ServiceUnavailable, SessionExpired)
        - Thread-safe: Neo4j driver handles concurrent session access

    Multi-tenancy:
        - If ``enable_multi_tenancy=True``, uses per-tenant databases
        - Otherwise, uses namespace prefixing within single database
    """

    def __init__(
        self,
        uri: str,
        user: str,
        password: str,
        *,
        default_namespace: str = "default",
        enable_multi_tenancy: bool = True,
        init_schema: bool = True,
        debug: bool = False
    ):
        super().__init__(debug=debug)
        self.uri = uri
        self.user = user
        self.password = password
        self.default_namespace = default_namespace

        self.enable_multi_tenancy = enable_multi_tenancy

        # Initialize tenant manager for multi-database support
        try:
            if enable_multi_tenancy:
                self.tenant_manager = Neo4jTenantManager(
                    uri, user, password, debug=debug)
                self.driver = self.tenant_manager.get_driver()
                self._log("Multi-tenancy enabled - using per-tenant databases")
            else:
                # Explicit pool configuration for production workloads
                self.driver = GraphDatabase.driver(
                    uri,
                    auth=(user, password),
                    max_connection_pool_size=50,  # Default is 100, tune based on load
                    connection_acquisition_timeout=60.0,  # seconds
                    connection_timeout=30.0,  # seconds
                )
                self.tenant_manager = None
                self._log("Multi-tenancy disabled - using namespace prefixing")

            # Verify connection immediately to fail fast
            self.driver.verify_connectivity()
            self._log(f"Connected to Neo4j at {uri}")

        except Exception as e:
            # If we cannot connect, and this class was instantiated, it's a critical error
            # because the user explicitly requested Neo4j.
            self.logger.error(f"Failed to connect to Neo4j at {uri}: {e}")
            raise ConnectionError(f"Could not connect to Neo4j GraphStore: {e}") from e

        # Schema initialization cache
        self._initialized_databases: set[str | None] = set()

        # Initialize schema if requested (can be disabled for testing)
        if init_schema:
            try:
                self._ensure_schema(None)
            except Exception as e:
                self.logger.warning(f"Could not initialize schema on startup: {e}")

        # Circuit breaker for Neo4j operations
        self.cb = CircuitBreaker("neo4j_store", failure_threshold=5, recovery_timeout=30)

    def _get_database_for_namespace(self, namespace: str) -> Optional[str]:
        """
        Get the database name for a namespace.

        If multi-tenancy is enabled, returns the tenant database name.
        Otherwise, returns None (uses default database).

        Args:
            namespace: Namespace identifier

        Returns:
            Database name or None
        """
        if not self.enable_multi_tenancy or not self.tenant_manager:
            return None

        # Extract tenant ID from namespace (format: tenant_id::namespace)
        if "::" in namespace:
            tenant_id = namespace.split("::")[0]
        else:
            tenant_id = namespace

        # Ensure database exists and return its name
        return self.tenant_manager.ensure_database(tenant_id)

    def _ensure_schema(self, database: Optional[str] = None) -> None:
        """Ensure schema constraints and indexes exist in the database."""
        if database in self._initialized_databases:
            return

        with self.driver.session(database=database) as session:
            session.run("""
                CREATE CONSTRAINT IF NOT EXISTS
                FOR (m:Memory) REQUIRE m.id IS UNIQUE
            """)
            session.run("""
                CREATE INDEX IF NOT EXISTS
                FOR (m:Memory) ON (m.embedding_type)
            """)
            session.run("""
                CREATE INDEX IF NOT EXISTS
                FOR (m:Memory) ON (m.session_id)
            """)
            session.run("""
                CREATE INDEX IF NOT EXISTS
                FOR (m:Memory) ON (m.namespace)
            """)

        self._initialized_databases.add(database)

    # ------------------------------------------------------------------ #
    # Episode persistence
    # ------------------------------------------------------------------ #
    def save_episode(self, ep: Episode) -> None:
        """Persist scene graph nodes and relations for an episode."""
        self.validate_episode(ep)
        G = ep.G
        metadata = ep.meta if isinstance(ep.meta, dict) else {}
        embedding_type = str(metadata.get("embedding_type", "default"))
        session_id = metadata.get("session_id")
        namespace = metadata.get("namespace") or self.default_namespace

        # Get tenant database
        database = self._get_database_for_namespace(namespace)

        # Ensure schema exists in tenant database
        if database:
            self._ensure_schema(database)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _save_op():
            with self.driver.session(database=database) as session:
                session.run(
                    """
                    MERGE (m:Memory {id:$eid})
                    ON CREATE SET m.created_at = datetime()
                    SET m.trust=$trust,
                        m.embedding_type=$embedding_type,
                        m.session_id=$session_id,
                        m.namespace=$namespace,
                        m.updated_at = datetime()
                    """,
                    eid=ep.id,
                    trust=float(ep.trust),
                    embedding_type=embedding_type,
                    session_id=session_id,
                    namespace=namespace,
                )
                for u, v, data in G.edges(data=True):
                    # Normalize relation key per GraphStore convention
                    rel = data.get("relation_type") or data.get("relation") or data.get("rel") or "RELATED_TO"
                    rel = str(rel).strip().upper()  # Neo4j relation types are uppercase
                    session.run(
                        """
                        MERGE (a:Entity {name: $u})
                        MERGE (b:Entity {name: $v})
                        WITH a, b, $props AS props
                        MATCH (m:Memory {id:$eid})
                        MERGE (m)-[:OBSERVED {namespace:$namespace}]->(a)
                        MERGE (m)-[:OBSERVED {namespace:$namespace}]->(b)
                        WITH a, b, props
                        MERGE (a)-[r:%s {episode:$eid, namespace:$namespace}]->(b)
                        SET r += props
                        """ % rel,
                        u=u,
                        v=v,
                        eid=ep.id,
                        namespace=namespace,
                        props={
                            **data,
                            "trust": ep.trust,
                            "tocc": ep.tocc.isoformat(),
                            "namespace": namespace,
                            "relation": rel.lower(),
                            "relation_type": rel.lower()
                        },
                    )

        self.cb.call(_save_op)
        self._log(f"Saved graph for episode {ep.id}")

    def save_many(self, episodes: Iterable[Episode]) -> None:
        episode_list = list(episodes)
        if not episode_list:
            return

        # Group episodes by namespace/tenant to batch per database
        from collections import defaultdict
        episodes_by_namespace = defaultdict(list)

        for ep in episode_list:
            metadata = ep.meta if isinstance(ep.meta, dict) else {}
            namespace = metadata.get("namespace") or self.default_namespace
            episodes_by_namespace[namespace].append(ep)

        # Process each namespace separately
        for namespace, ns_episodes in episodes_by_namespace.items():
            database = self._get_database_for_namespace(namespace)
            if database:
                self._ensure_schema(database)

            # Prepare batch data for this namespace
            memory_nodes = []
            relationships = []

            for ep in ns_episodes:
                self.validate_episode(ep)
                metadata = ep.meta if isinstance(ep.meta, dict) else {}

                # Memory node data
                memory_nodes.append({
                    "eid": ep.id,
                    "trust": float(ep.trust),
                    "embedding_type": str(metadata.get("embedding_type", "default")),
                    "session_id": metadata.get("session_id"),
                    "namespace": namespace,
                })

                # Relationships
                for u, v, data in ep.G.edges(data=True):
                    # Normalize relation key per GraphStore convention
                    rel_type = data.get("relation_type") or data.get("relation") or data.get("rel") or "RELATED_TO"
                    rel_type = str(rel_type).strip().upper()
                    rel_normalized = rel_type.lower()
                    props = {
                        **data,
                        "trust": ep.trust,
                        "tocc": ep.tocc.isoformat(),
                        "namespace": namespace,
                        "relation": rel_normalized,
                        "relation_type": rel_normalized
                    }
                    props.pop("rel", None)

                    relationships.append({
                        "eid": ep.id,
                        "u": u,
                        "v": v,
                        "rel_type": rel_type,
                        "namespace": namespace,
                        "props": props
                    })

            @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
            def _save_many_op():
                with self.driver.session(database=database) as session:
                    # 1. Batch create Memory nodes
                    session.run(
                        """
                    UNWIND $batch AS item
                    MERGE (m:Memory {id: item.eid})
                    ON CREATE SET m.created_at = datetime()
                    SET m.trust = item.trust,
                        m.embedding_type = item.embedding_type,
                        m.session_id = item.session_id,
                        m.namespace = item.namespace,
                        m.updated_at = datetime()
                    """,
                        batch=memory_nodes
                    )

                    # 2. Batch create Entities and Relationships
                    from collections import defaultdict
                    rels_by_type = defaultdict(list)
                    for rel in relationships:
                        rels_by_type[rel["rel_type"]].append(rel)

                    for rel_type, batch in rels_by_type.items():
                        query = f"""
                        UNWIND $batch AS item
                        MERGE (a:Entity {{name: item.u}})
                        MERGE (b:Entity {{name: item.v}})
                        WITH a, b, item
                        MATCH (m:Memory {{id: item.eid}})
                        MERGE (m)-[:OBSERVED {{namespace: item.namespace}}]->(a)
                        MERGE (m)-[:OBSERVED {{namespace: item.namespace}}]->(b)
                        WITH a, b, item
                        MERGE (a)-[r:{rel_type} {{episode: item.eid, namespace: item.namespace}}]->(b)
                        SET r += item.props
                        """
                        session.run(query, batch=batch)

            self.cb.call(_save_many_op)

        self._log(f"Saved {len(episode_list)} episodes to Neo4j (batched)")

    # ------------------------------------------------------------------ #
    # Loading / Querying
    # ------------------------------------------------------------------ #
    def load_all(self, *, namespace: Optional[str] = None) -> list[Episode]:
        """Load all graphs back into Episode objects (without embeddings)."""
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _load_all_op():
            with self.driver.session(database=database) as session:
                result = session.run(
                    """
                    MATCH (m:Memory {namespace:$namespace})-[:OBSERVED {namespace:$namespace}]->(a)
                    MATCH (m)-[r {namespace:$namespace}]->(b)
                    RETURN a.name AS src, b.name AS dst, type(r) AS rel, r.episode AS eid, r
                    """,
                    namespace=graph_ns,
                )
                return result.data()

        records = self.cb.call(_load_all_op)

        episodes: dict[str, Episode] = {}
        for rec in records:
            eid = rec["eid"]
            if not eid:
                continue
            if eid not in episodes:
                G = nx.DiGraph()
                default_vec = np.zeros(1, dtype=np.float32)
                default_time = datetime.fromtimestamp(0)
                trust = float(rec["r"].get("trust", 1.0))
                episodes[eid] = Episode(
                    id=eid,
                    e=default_vec,
                    G=G,
                    tocc=default_time,
                    meta={},
                    parent=None,
                    trust=trust,
                )
            G = episodes[eid].G
            G.add_edge(
                rec["src"], rec["dst"],
                relation=rec["rel"],
                **{k: v for k, v in rec["r"].items() if k not in {"episode", "rel"}}
            )
        self._log(f"Loaded {len(episodes)} episode graphs from Neo4j")
        return list(episodes.values())

    def load_by_id(self, episode_id: str, *, namespace: Optional[str] = None) -> Optional[Episode]:
        """Load a single episode graph by ID."""
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _load_by_id_op():
            with self.driver.session(database=database) as session:
                result = session.run(
                    """
                    MATCH (a)-[r {episode:$eid, namespace:$namespace}]->(b)
                    RETURN a.name AS src, b.name AS dst, type(r) AS rel, r
                    """,
                    eid=episode_id,
                    namespace=graph_ns,
                )
                return result.data()

        records = self.cb.call(_load_by_id_op)

        if not records:
            return None

        G = nx.DiGraph()
        for rec in records:
            G.add_edge(
                rec["src"], rec["dst"],
                relation=rec["rel"],
                **{k: v for k, v in rec["r"].items() if k not in {"episode", "rel"}}
            )
        default_vec = np.zeros(1, dtype=np.float32)
        default_time = datetime.fromtimestamp(0)
        return Episode(
            id=episode_id,
            e=default_vec,
            G=G,
            tocc=default_time,
            meta={},
            parent=None,
            trust=1.0,
        )

    # ------------------------------------------------------------------ #
    # Updating & Deleting
    # ------------------------------------------------------------------ #
    def update_trust(self, episode_id: str, trust: float, *, namespace: Optional[str] = None) -> None:
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _update_trust_op():
            with self.driver.session(database=database) as session:
                session.run(
                    "MATCH ()-[r {episode:$eid, namespace:$namespace}]->() SET r.trust=$trust",
                    eid=episode_id,
                    trust=trust,
                    namespace=graph_ns,
                )
        self.cb.call(_update_trust_op)
        self._log(f"Updated trust for episode {episode_id} → {trust:.3f}")

    def remove(self, episode_id: str, *, namespace: Optional[str] = None) -> None:
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _remove_op():
            with self.driver.session(database=database) as session:
                session.run(
                    "MATCH ()-[r {episode:$eid, namespace:$namespace}]->() DELETE r",
                    eid=episode_id,
                    namespace=graph_ns,
                )
        self.cb.call(_remove_op)
        self._log(f"Removed all relations for episode {episode_id}")

    def remove_many(self, episode_ids: Iterable[str], *, namespace: Optional[str] = None) -> None:
        ids = list(episode_ids)
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _remove_many_op():
            with self.driver.session(database=database) as session:
                session.run(
                    "UNWIND $ids AS eid MATCH ()-[r {episode:eid, namespace:$namespace}]->() DELETE r",
                    ids=ids,
                    namespace=graph_ns,
                )
        self.cb.call(_remove_many_op)
        self._log(f"Removed {len(ids)} episode graphs")

    def clear(self) -> None:
        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _clear_op():
            # Clear all databases when multi-tenancy is enabled
            if self.enable_multi_tenancy and self.tenant_manager:
                tenant_dbs = self.tenant_manager.list_tenant_databases()
                for db_name in tenant_dbs.values():
                    with self.driver.session(database=db_name) as session:
                        session.run("MATCH (n) DETACH DELETE n")
                self._log(f"Cleared {len(tenant_dbs)} tenant databases")
            else:
                # Clear default database
                with self.driver.session() as session:
                    session.run("MATCH (n) DETACH DELETE n")
                self._log("Cleared default database")

        self.cb.call(_clear_op)
        self._log("Cleared entire Neo4j graph database")

    # ------------------------------------------------------------------ #
    # Symbolic Queries (GraphStore Implementation)
    # ------------------------------------------------------------------ #
    def query_relations(self, entity: str, depth: int = 1, *, namespace: Optional[str] = None) -> list[dict[str, Any]]:
        """Return all outgoing relations from a given entity node."""
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _query_relations_op():
            with self.driver.session(database=database) as session:
                # Assuming depth 1 for basic query_relations backward compat, or use variable length path
                # GraphStore defines signature as `query_relations(self, entity: str, depth: int = 1)`
                # We can implement traversal.

                # Note: The original returned tuple[str, str, str].
                # Newer GraphStore contract returns list[dict].
                result = session.run(
                    f"""
                    MATCH (a {{name:$entity}})-[r*1..{depth}]->(b)
                    RETURN a.name AS src, [rel in r | type(rel)] AS relations, b.name AS dst,
                           [rel in r | rel.trust] as trusts,
                           [rel in r | properties(rel)] as props
                    LIMIT 200
                    """,
                    entity=entity,
                )

                data = []
                for rec in result:
                    # Return first relation for compatibility with InMemoryGraphStore schema
                    first_rel = rec["relations"][0] if rec["relations"] else "RELATED_TO"
                    data.append({
                        "src": rec["src"],
                        "dst": rec["dst"],
                        "relation": first_rel.lower(),  # Canonical lowercase
                        "trust": rec["trusts"][0] if rec["trusts"] else 1.0,
                        "depth": len(rec["relations"]),  # Path length = depth
                        "properties": rec["props"][0] if rec["props"] else {}
                    })
                return data

        data = self.cb.call(_query_relations_op)
        self._log(f"Queried relations for entity '{entity}' depth={depth}")
        return data

    def search_related(
        self,
        episode_id: str,
        depth: int = 2,
        limit: int = 50,
        *,
        namespace: Optional[str] = None
    ) -> list[dict]:
        """Expand entities observed in an episode with bounded traversal depth."""
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        depth = max(1, min(int(depth), 4))
        limit = max(1, min(int(limit), 200))

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _search_related_op():
            with self.driver.session(database=database) as session:
                result = session.run(
                    f"""
                    MATCH (m:Memory {{id:$eid}})-[:OBSERVED]->(seed:Entity)
                    MATCH path=(seed)-[r*1..{depth}]->(neighbor:Entity)
                    RETURN seed.name AS source,
                           neighbor.name AS target,
                           [rel IN relationships(path) | type(rel)] AS relations,
                           [rel IN relationships(path) | rel.trust] AS trusts,
                           [rel IN relationships(path) | properties(rel)] AS props
                    LIMIT $limit
                    """,
                    eid=episode_id,
                    limit=limit,
                )

                data = []
                for rec in result:
                    data.append({
                        "source": rec["source"],
                        "target": rec["target"],
                        "relations": rec["relations"],
                        "trust": rec["trusts"][0] if rec["trusts"] else 1.0,
                        "properties": rec["props"]
                    })
                return data

        data = self.cb.call(_search_related_op)
        self._log(f"Expanded episode {episode_id} to {len(data)} related entities")
        return data

    def get_graph_stats(self, *, namespace: Optional[str] = None) -> dict[str, Any]:
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        @with_retry(exceptions=(ServiceUnavailable, SessionExpired, TransientError))
        def _stats():
            with self.driver.session(database=database) as session:
                res = session.run("MATCH (n) RETURN count(n) as node_count")
                nodes = res.single()["node_count"]

                res2 = session.run("MATCH ()-[r]->() RETURN count(r) as edge_count")
                edges = res2.single()["edge_count"]
                return {
                    "enabled": True,
                    "node_count": nodes,
                    "edge_count": edges,
                    "backend": "neo4j",
                    "type": "Neo4jStore"
                }

        return self.cb.call(_stats)

    # ------------------------------------------------------------------ #
    # Query Profiling for Performance Analysis
    # ------------------------------------------------------------------ #
    def profile_query(self, cypher: str, params: dict | None = None, *, namespace: str | None = None) -> dict[str, Any]:
        """
        Profile a Cypher query using Neo4j's PROFILE command.

        Returns execution plan with timing and resource usage.
        Useful for identifying slow queries and optimizing indexes.

        Args:
            cypher: The Cypher query to profile (without PROFILE prefix)
            params: Query parameters
            namespace: Optional namespace for database selection

        Returns:
            Dict with profile results including dbHits, rows, and timing
        """
        graph_ns = namespace or self.default_namespace
        database = self._get_database_for_namespace(graph_ns)

        with self.driver.session(database=database) as session:
            result = session.run(f"PROFILE {cypher}", params or {})
            summary = result.consume()

            profile_data = {
                "query": cypher,
                "result_available_after_ms": summary.result_available_after,
                "result_consumed_after_ms": summary.result_consumed_after,
                "counters": {
                    "nodes_created": summary.counters.nodes_created,
                    "relationships_created": summary.counters.relationships_created,
                    "properties_set": summary.counters.properties_set,
                },
                "plan": str(summary.plan) if summary.plan else None,
                "profile": str(summary.profile) if summary.profile else None,
            }
            return profile_data

    # ------------------------------------------------------------------ #
    # Context Management
    # ------------------------------------------------------------------ #
    def close(self) -> None:
        self.driver.close()
        self._log("Neo4j driver closed")
