"""
graph_store.py
--------------
Abstract base class for graph-specific persistence in HIEROMEM.
Defines the contract for storing and querying symbolic structures (scene graphs).

This interface decouples the core memory system from specific graph databases
(Neo4j, NetworkX, RDF, etc.), ensuring that graph capabilities are explicit
and testable.

Author: Antigravity
"""

from __future__ import annotations
from abc import abstractmethod
from typing import Any, Optional

from hieromem_core.persistence.memory_store import MemoryStore


class GraphStore(MemoryStore):
    """
    Abstract interface for graph-capable storage backends.

    Inherits from MemoryStore to support basic persistence (save/load episode),
    but adds methods specific to symbolic reasoning and graph traversal.

    Edge Attribute Convention:
        All implementations MUST store edges with the following canonical attributes:
        - ``relation``: The normalized relation type (lowercase, stripped)
        - ``relation_type``: Alias for ``relation`` (for compatibility)
        - ``namespace``: The namespace this edge belongs to
        - ``episode``: The episode ID that created this edge
        - ``trust``: Trust score for the edge (default 1.0)

        Implementations should normalize incoming edges at ingestion time:
        - Read from ``relation_type``, ``relation``, or ``rel``
        - Store as both ``relation`` and ``relation_type``
        - Normalize value: ``str(rel).strip().lower()``
    """

    @abstractmethod
    def query_relations(self, entity: str, depth: int = 1, *, namespace: Optional[str] = None) -> list[dict[str, Any]]:
        """
        Retrieve relationships originating from or involving the given entity.

        Args:
            entity: The entity name/ID to query.
            depth: Traversal depth (default 1).
            namespace: Optional tenant namespace.

        Returns:
            List of relationship records (source, target, relation type, properties).
        """

    @abstractmethod
    def search_related(
        self,
        episode_id: str,
        depth: int = 2,
        limit: int = 50,
        *,
        namespace: Optional[str] = None
    ) -> list[dict[str, Any]]:
        """
        Expand entities observed in a specific episode to find related context.

        Args:
            episode_id: The episode to use as a seed.
            depth: Max hops to traverse.
            limit: Max number of results.
            namespace: Optional tenant namespace.

        Returns:
            List of paths or subgraphs related to the episode's entities.
        """

    @abstractmethod
    def get_graph_stats(self, *, namespace: Optional[str] = None) -> dict[str, Any]:
        """
        Return statistics about the stored graph (node count, edge count, density).

        Returns:
            Dict containing metrics like 'node_count', 'edge_count'.
        """

    # ------------------------------------------------------------------ #
    # Relation Normalization Helper
    # ------------------------------------------------------------------ #
    @staticmethod
    def normalize_relation_key(edge_attrs: dict[str, Any]) -> str:
        """
        Extract and normalize the relation key from edge attributes.

        Implementations SHOULD use this method during ingestion to ensure
        consistent relation keys across the system.

        Reads from: ``relation_type``, ``relation``, or ``rel`` (in priority order)
        Returns: Lowercase, stripped relation string

        Args:
            edge_attrs: Dictionary of edge attributes

        Returns:
            Normalized relation string (e.g., "related_to")

        Example:
            >>> attrs = {"relation_type": " OWNS ", "weight": 0.5}
            >>> GraphStore.normalize_relation_key(attrs)
            'owns'
        """
        rel = edge_attrs.get("relation_type") or edge_attrs.get("relation")
        rel = rel or edge_attrs.get("rel") or "RELATED_TO"
        return str(rel).strip().lower()

    @staticmethod
    def apply_canonical_relation(edge_attrs: dict[str, Any]) -> dict[str, Any]:
        """
        Normalize edge attributes and ensure both canonical relation keys are set.

        This modifies the edge_attrs dict in place and returns it.

        Args:
            edge_attrs: Dictionary of edge attributes (modified in place)

        Returns:
            The modified edge_attrs dict with 'relation' and 'relation_type' set
        """
        rel = GraphStore.normalize_relation_key(edge_attrs)
        edge_attrs["relation"] = rel
        edge_attrs["relation_type"] = rel
        # Remove non-canonical keys to avoid confusion
        edge_attrs.pop("rel", None)
        return edge_attrs

    # ------------------------------------------------------------------ #
    # 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

        Note:
            Default implementation is a no-op. Subclasses should override
            to persist structural edges.
        """
        pass  # Default no-op for stores that don't support structural edges

    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)

        Note:
            Default returns empty list. Subclasses should override.
        """
        return []

    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)

        Note:
            Default returns empty list. Subclasses should override.
        """
        return []

    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.

        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)

        Note:
            Default returns empty list. Subclasses should override.
        """
        return []
