"""
in_memory.py
------------
In-Memory NetworkX implementation of GraphBackend.
"""

from typing import Any, Dict, List, Optional, Tuple
import networkx as nx

from hieromem_core.backends.interface import GraphBackend


class InMemoryGraphBackend(GraphBackend):
    def __init__(self):
        self.graph = nx.MultiDiGraph()

    def add_triples(self, triples: List[Tuple[str, str, str, Dict[str, Any]]]) -> None:
        for u, rel, v, props in triples:
            self.graph.add_edge(u, v, key=rel, relation=rel, **props)

    def search_triples(
        self,
        head: Optional[str] = None,
        relation: Optional[str] = None,
        tail: Optional[str] = None,
    ) -> List[Tuple[str, str, str, Dict[str, Any]]]:
        results = []
        # Optimization: use graph indices if needed. NetworkX is slow for full scan.
        # If head is known:
        if head and head in self.graph:
            # Outgoing edges
            for v, key, data in self.graph.edges(head, keys=True, data=True):
                r = data.get("relation")
                if relation and r != relation:
                    continue
                if tail and v != tail:
                    continue
                results.append((head, r, v, data))

        # If head is None but tail is known (NetworkX reverse lookup is slow on MultiDiGraph without pre-index)
        elif tail and not head:
            for u, v, key, data in self.graph.in_edges(tail, keys=True, data=True):
                r = data.get("relation")
                if relation and r != relation:
                    continue
                results.append((u, r, tail, data))

        else:
            # Full scan (slow)
            for u, v, key, data in self.graph.edges(keys=True, data=True):
                if head and u != head:
                    continue
                if tail and v != tail:
                    continue
                r = data.get("relation")
                if relation and r != relation:
                    continue
                results.append((u, r, v, data))

        return results

    def get_neighbors(
        self, node: str, depth: int = 1
    ) -> List[Tuple[str, str, str, Dict[str, Any]]]:
        if node not in self.graph:
            return []
        # Simple depth-1 implementation
        results = []
        for v, key, data in self.graph.edges(node, keys=True, data=True):
            results.append((node, data.get("relation"), v, data))
        return results

    def delete_triples(self, triples: List[Tuple[str, str, str]]) -> None:
        for u, r, v in triples:
            # Need to find key for (u, v) with relation r
            to_remove = []
            if self.graph.has_edge(u, v):
                for key, data in self.graph[u][v].items():
                    if data.get("relation") == r:
                        to_remove.append(key)

                for key in to_remove:
                    self.graph.remove_edge(u, v, key=key)

    def clear(self) -> None:
        self.graph.clear()
