"""
in_memory_graph_store.py
------------------------
In-memory implementation of GraphStore using NetworkX.
Suitable for testing, CI protocols, and lightweight deployments where
persistence beyond the process lifetime is not required (or handled via simple serialization).

Author: Antigravity
"""

from __future__ import annotations
import networkx as nx
import threading
from collections import deque, defaultdict
from typing import Any, Iterable, Optional, List
import uuid

from hieromem_core.episode import Episode
from hieromem_core.edge_types import (
    TURN_FOLLOWS, SPOKEN_BY, IN_SESSION,
    SPEAKER_NODE_PREFIX, SESSION_NODE_PREFIX,
)
from .graph_store import GraphStore


class InMemoryGraphStore(GraphStore):
    """
    GraphStore implementation backed by a single NetworkX DiGraph.

    Acts as a "Global Graph" where all episodes are merged into one structure,
    enabling multi-hop reasoning across the full memory.

    WARNING: This implementation is NOT thread-safe. NetworkX graphs do not
    support concurrent modification. For multi-threaded deployments, use
    Neo4jStore which handles concurrency internally, or add external locking.

    Attributes:
        _G: NetworkX MultiDiGraph containing all symbolic relationships.
        _episodes: Dict mapping episode ID to Episode objects (for load_all).
        graph_store: Self-reference for Memory compatibility.
    """

    def __init__(self, *, debug: bool = False):
        super().__init__(debug=debug)
        self._G: nx.MultiDiGraph = nx.MultiDiGraph()
        self._episodes: dict[str, Episode] = {}
        # Index for O(1) edge lookup by episode ID (performance fix for update_trust)
        self._edges_by_episode: dict[str, list[tuple[Any, Any, str]]] = defaultdict(list)
        # Threading lock for concurrent safety (single-threaded by default)
        self._lock = threading.RLock()
        # Set graph_store to self so Memory can find it via store.graph_store
        self.graph_store = self

    # ------------------------------------------------------------------ #
    # MemoryStore Interface (episodes as primary)
    # ------------------------------------------------------------------ #
    def save_episode(self, ep: Episode, *, namespace: Optional[str] = None) -> None:
        with self._lock:
            self._episodes[ep.id] = ep

            # Merge symbolic graph (G) into global graph
            if not hasattr(ep, "G") or ep.G is None:
                return

            # Normalize namespace: ensure it's a non-empty string (ephemeral namespace fix)
            ns = namespace or ep.meta.get("namespace") or "default"
            ns = str(ns).strip() or "default"
            # Write normalized namespace back to ep.meta for consistency
            if isinstance(ep.meta, dict):
                ep.meta["namespace"] = ns

            for u, v, edge_attrs in ep.G.edges(data=True):
                edge_attrs = dict(edge_attrs)
                edge_attrs["episode"] = ep.id
                edge_attrs["namespace"] = ns

                # Normalize relation key - store as both relation_type AND relation
                rel_type = edge_attrs.get("relation_type") or edge_attrs.pop("relation", None)
                rel_type = rel_type or edge_attrs.pop("rel", None) or "RELATED_TO"
                rel_type = str(rel_type).strip().lower()
                edge_attrs["relation_type"] = rel_type
                edge_attrs["relation"] = rel_type

                # Create nodes BEFORE adding edge (fix: node update order)
                # NetworkX implicitly creates nodes, but we need to set namespaces first
                for node in (u, v):
                    if node not in self._G:
                        self._G.add_node(node, namespaces={ns})
                    else:
                        # Add this namespace to existing node's namespace set
                        existing_ns = self._G.nodes[node].get("namespaces")
                        if isinstance(existing_ns, set):
                            existing_ns.add(ns)
                        else:
                            self._G.nodes[node]["namespaces"] = {ns}

                # Use unique key to avoid overwriting parallel edges
                key = f"{ep.id}_{rel_type}_{uuid.uuid4().hex[:8]}"
                self._G.add_edge(u, v, key=key, **edge_attrs)

                # Register edge in index for O(1) lookup by episode (performance fix)
                self._edges_by_episode[ep.id].append((u, v, key))

            self._log(f"Merged episode {ep.id} into in-memory graph")

    def save_many(self, episodes: Iterable[Episode]) -> None:
        for ep in episodes:
            self.save_episode(ep)

    def load_all(self, *, namespace: Optional[str] = None) -> list[Episode]:
        all_eps = [ep.clone(new_id=False) for ep in self._episodes.values()]
        if namespace:
            return [
                ep for ep in all_eps
                if (ep.meta.get("namespace") == namespace) or (namespace == "default" and "namespace" not in ep.meta)
            ]
        return all_eps

    def load_by_id(self, episode_id: str, *, namespace: Optional[str] = None) -> Optional[Episode]:
        ep = self._episodes.get(episode_id)
        if not ep:
            return None

        if namespace:
            ep_ns = ep.meta.get("namespace", "default")
            if ep_ns != namespace:
                return None

        return ep.clone(new_id=False)

    def update_trust(self, episode_id: str, trust: float, *, namespace: Optional[str] = None) -> None:
        with self._lock:
            # Update in storage dict
            if episode_id in self._episodes:
                self._episodes[episode_id].trust = float(trust)

            # Use edge index for O(1) lookup instead of O(E) scan (performance fix)
            if episode_id in self._edges_by_episode:
                for u, v, k in self._edges_by_episode[episode_id]:
                    if (u, v, k) in self._G.edges(keys=True):
                        data = self._G[u][v][k]
                        if namespace and data.get("namespace") != namespace:
                            continue
                        self._G[u][v][k]["trust"] = float(trust)

    def remove(self, episode_id: str, *, namespace: Optional[str] = None) -> None:
        with self._lock:
            # Remove from dict
            if episode_id in self._episodes:
                if namespace:
                    ep = self._episodes[episode_id]
                    ep_ns = ep.meta.get("namespace", "default")
                    if ep_ns != namespace:
                        return
                del self._episodes[episode_id]

            # Use edge index for efficient removal
            edges_to_remove = []
            if episode_id in self._edges_by_episode:
                for u, v, k in self._edges_by_episode[episode_id]:
                    if (u, v, k) in self._G.edges(keys=True):
                        data = self._G[u][v][k]
                        if namespace and data.get("namespace") != namespace:
                            continue
                        edges_to_remove.append((u, v, k))

            for u, v, k in edges_to_remove:
                self._G.remove_edge(u, v, k)

            # Clean up edge index
            if episode_id in self._edges_by_episode:
                del self._edges_by_episode[episode_id]

            # Cleanup isolated nodes
            nodes_to_check = {u for u, _, _ in edges_to_remove} | {v for _, v, _ in edges_to_remove}
            for n in nodes_to_check:
                if n in self._G and self._G.degree(n) == 0:
                    self._G.remove_node(n)

    def remove_many(self, episode_ids: Iterable[str], *, namespace: Optional[str] = None) -> None:
        for eid in episode_ids:
            self.remove(eid, namespace=namespace)

    def clear(self) -> None:
        with self._lock:
            self._episodes.clear()
            self._G.clear()
            self._edges_by_episode.clear()

    # ------------------------------------------------------------------ #
    # GraphStore Implementation (Symbolic)
    # ------------------------------------------------------------------ #
    def query_relations(self, entity: str, depth: int = 1, *, namespace: Optional[str] = None) -> list[dict[str, Any]]:
        """
        Query relations from an entity with proper depth support (Bug #3 fix).
        Uses BFS to traverse up to `depth` hops.
        """
        if entity not in self._G:
            return []

        results = []
        visited = {entity}
        queue: deque = deque([(entity, 0)])

        while queue:
            current_node, current_depth = queue.popleft()
            if current_depth >= depth:
                continue

            # Find outgoing edges from current node
            if current_node in self._G:
                for _, v, k, data in self._G.out_edges(current_node, data=True, keys=True):
                    # Namespace filter on edge
                    if namespace and data.get("namespace") != namespace:
                        continue

                    # Check destination node namespace before including in results
                    if namespace:
                        node_namespaces = self._G.nodes[v].get("namespaces", set())
                        if namespace not in node_namespaces:
                            continue

                    excluded_keys = ("episode", "namespace", "trust", "relation_type")
                    results.append({
                        "src": current_node,
                        "dst": v,
                        "relation": data.get("relation_type", "RELATED_TO"),
                        "trust": data.get("trust", 1.0),
                        "depth": current_depth + 1,
                        "properties": {k: v for k, v in data.items() if k not in excluded_keys}
                    })

                    # Only traverse to unvisited nodes in this namespace
                    if v not in visited:
                        if not namespace or namespace in self._G.nodes[v].get("namespaces", set()):
                            visited.add(v)
                            queue.append((v, current_depth + 1))

        return results

    def search_related(
        self,
        episode_id: str,
        depth: int = 2,
        limit: int = 50,
        *,
        namespace: Optional[str] = None
    ) -> list[dict[str, Any]]:
        """
        BFS traversal from seed nodes with proper deque and early termination (Bug #4 fix).
        """
        # 1. Find seed entities from this episode
        seed_nodes = set()
        for u, v, k, data in self._G.edges(data=True, keys=True):
            if data.get("episode") == episode_id:
                if namespace and data.get("namespace") != namespace:
                    continue
                seed_nodes.add(u)
                seed_nodes.add(v)

        if not seed_nodes:
            return []

        # 2. BFS traversal using deque for O(1) popleft (Bug #4 fix)
        visited = set(seed_nodes)
        queue: deque = deque([(n, 0) for n in seed_nodes])
        collected_edges = []

        while queue:
            # Early termination if limit reached (Bug #4 fix)
            if len(collected_edges) >= limit:
                break

            current_node, current_depth = queue.popleft()
            if current_depth >= depth:
                continue

            # Traversal
            if current_node in self._G:
                for _, neighbor, k, data in self._G.out_edges(current_node, data=True, keys=True):
                    # Namespace check
                    if namespace and data.get("namespace") != namespace:
                        continue

                    # Capture edge
                    edge_record = {
                        "source": current_node,
                        "target": neighbor,
                        "relations": [data.get("relation_type", "RELATED_TO")],
                        "trust": data.get("trust", 1.0)
                    }
                    collected_edges.append(edge_record)

                    # Early termination check
                    if len(collected_edges) >= limit:
                        return collected_edges

                    # Check that neighbor node belongs to namespace (namespace leak fix)
                    if namespace:
                        neighbor_namespaces = self._G.nodes[neighbor].get("namespaces", set())
                        if namespace not in neighbor_namespaces:
                            continue

                    if neighbor not in visited:
                        visited.add(neighbor)
                        queue.append((neighbor, current_depth + 1))

        return collected_edges

    def get_graph_stats(self, *, namespace: Optional[str] = None) -> dict[str, Any]:
        """Return graph statistics with consistent enabled flag."""
        # If namespace requested, count only nodes/edges belonging to it
        if namespace:
            edge_count = sum(
                1 for _, _, data in self._G.edges(data=True)
                if data.get("namespace") == namespace
            )
            node_count = sum(
                1 for n in self._G.nodes
                if namespace in self._G.nodes[n].get("namespaces", set())
            )
        else:
            node_count = self._G.number_of_nodes()
            edge_count = self._G.number_of_edges()

        return {
            "enabled": True,
            "node_count": node_count,
            "edge_count": edge_count,
            "backend": "networkx",
            "type": "InMemoryGraphStore"
        }

    # ------------------------------------------------------------------ #
    # Structural Edge Methods (TURN_FOLLOWS, SPOKEN_BY, IN_SESSION)
    # ------------------------------------------------------------------ #
    def add_structural_edge(
        self,
        source: str,
        target: str,
        edge_attrs: dict[str, Any],
        *,
        namespace: Optional[str] = None
    ) -> None:
        """
        Add a structural edge (TURN_FOLLOWS, SPOKEN_BY, IN_SESSION).

        These edges create conversational scaffolding for temporal,
        speaker-based, and session-based queries.

        Args:
            source: Source node ID (episode ID or other node)
            target: Target node ID (episode ID, speaker node, or session node)
            edge_attrs: Edge attributes including relation_type, weight, etc.
            namespace: Optional namespace for multi-tenant
        """
        with self._lock:
            ns = namespace or edge_attrs.get("namespace") or "default"
            edge_attrs = dict(edge_attrs)
            edge_attrs["namespace"] = ns

            # Create nodes if they don't exist, with appropriate node_type
            for node in (source, target):
                if node not in self._G:
                    # Determine node_type based on prefix
                    node_type = "episode"
                    if node.startswith(SPEAKER_NODE_PREFIX):
                        node_type = "speaker"
                    elif node.startswith(SESSION_NODE_PREFIX):
                        node_type = "session"

                    self._G.add_node(node, namespaces={ns}, node_type=node_type)
                else:
                    existing_ns = self._G.nodes[node].get("namespaces")
                    if isinstance(existing_ns, set):
                        existing_ns.add(ns)
                    else:
                        self._G.nodes[node]["namespaces"] = {ns}

            # Use unique key to support parallel edges
            rel_type = edge_attrs.get("relation_type", "structural")
            key = f"{source}_{target}_{rel_type}_{uuid.uuid4().hex[:8]}"
            self._G.add_edge(source, target, key=key, **edge_attrs)

            self._log(f"Added structural edge: {source} -[{rel_type}]-> {target}")

    def query_turns_after(
        self,
        episode_id: str,
        limit: int = 10,
        *,
        namespace: Optional[str] = None
    ) -> List[str]:
        """
        Get episode IDs that follow the given episode temporally.

        Traverses TURN_FOLLOWS edges from the given episode.

        Args:
            episode_id: Starting episode ID
            limit: Maximum number of following episodes to return
            namespace: Optional namespace filter

        Returns:
            List of episode IDs in temporal order (closest first)
        """
        if episode_id not in self._G:
            self._log(f"query_turns_after: Episode {episode_id} not found in graph")
            return []

        result = []
        visited = {episode_id}
        current = episode_id

        while len(result) < limit:
            # Find outgoing TURN_FOLLOWS edges
            found_next = False
            if current in self._G:
                for _, target, k, data in self._G.out_edges(current, data=True, keys=True):
                    rel = data.get("relation_type", data.get("relation", ""))
                    if rel != TURN_FOLLOWS:
                        continue
                    if namespace and data.get("namespace") != namespace:
                        continue
                    if target in visited:
                        continue

                    visited.add(target)
                    result.append(target)
                    current = target
                    found_next = True
                    break

            if not found_next:
                break

        return result

    def query_by_speaker(
        self,
        speaker: str,
        limit: int = 50,
        *,
        namespace: Optional[str] = None
    ) -> List[str]:
        """
        Get episode IDs spoken by a specific speaker.

        Finds episodes connected to the speaker node via SPOKEN_BY edges.

        Args:
            speaker: Speaker name/ID to search for
            limit: Maximum number of episodes to return
            namespace: Optional namespace filter

        Returns:
            List of episode IDs (ordered by weight, highest first)
        """
        speaker_node = f"{SPEAKER_NODE_PREFIX}{speaker.strip()}"

        if speaker_node not in self._G:
            self._log(f"query_by_speaker: Speaker '{speaker}' not found in graph")
            return []

        episodes_with_weights = []

        # Find incoming SPOKEN_BY edges to the speaker node
        for source, _, k, data in self._G.in_edges(speaker_node, data=True, keys=True):
            rel = data.get("relation_type", data.get("relation", ""))
            if rel != SPOKEN_BY:
                continue
            if namespace and data.get("namespace") != namespace:
                continue

            weight = data.get("weight", 1.0)
            episodes_with_weights.append((source, weight))

        # Sort by weight (higher weight = more relevant)
        episodes_with_weights.sort(key=lambda x: x[1], reverse=True)

        return [ep_id for ep_id, _ in episodes_with_weights[:limit]]

    def query_by_session(
        self,
        session_id: str,
        limit: int = 100,
        *,
        namespace: Optional[str] = None
    ) -> List[str]:
        """
        Get episode IDs in a specific session.

        Finds episodes connected to the session node via IN_SESSION edges.
        Episodes are ordered by turn_index if available.

        Args:
            session_id: Session identifier to search for
            limit: Maximum number of episodes to return
            namespace: Optional namespace filter

        Returns:
            List of episode IDs (ordered by turn_index if available)
        """
        session_node = f"{SESSION_NODE_PREFIX}{session_id.strip()}"

        if session_node not in self._G:
            self._log(f"query_by_session: Session '{session_id}' not found in graph")
            return []

        episodes_with_index = []

        # Find incoming IN_SESSION edges to the session node
        for source, _, k, data in self._G.in_edges(session_node, data=True, keys=True):
            rel = data.get("relation_type", data.get("relation", ""))
            if rel != IN_SESSION:
                continue
            if namespace and data.get("namespace") != namespace:
                continue

            turn_index = data.get("turn_index", float("inf"))
            episodes_with_index.append((source, turn_index))

        # Sort by turn_index (chronological order within session)
        episodes_with_index.sort(key=lambda x: x[1])

        return [ep_id for ep_id, _ in episodes_with_index[:limit]]

    def get_speakers(self, *, namespace: Optional[str] = None) -> List[str]:
        """
        Get all speaker names in the graph.

        Args:
            namespace: Optional namespace filter

        Returns:
            List of speaker names
        """
        speakers = []
        for node in self._G.nodes:
            if node.startswith(SPEAKER_NODE_PREFIX):
                if namespace:
                    node_ns = self._G.nodes[node].get("namespaces", set())
                    if namespace not in node_ns:
                        continue
                speakers.append(node[len(SPEAKER_NODE_PREFIX):])
        return speakers

    def get_sessions(self, *, namespace: Optional[str] = None) -> List[str]:
        """
        Get all session IDs in the graph.

        Args:
            namespace: Optional namespace filter

        Returns:
            List of session IDs
        """
        sessions = []
        for node in self._G.nodes:
            if node.startswith(SESSION_NODE_PREFIX):
                if namespace:
                    node_ns = self._G.nodes[node].get("namespaces", set())
                    if namespace not in node_ns:
                        continue
                sessions.append(node[len(SESSION_NODE_PREFIX):])
        return sessions
