"""
memory_store.py
----------------
Abstract interface for persistence backends in HIEROMEM.

Every concrete store (FileStore, PgVectorStore, Neo4jStore, FAISSIndex)
must inherit from this class and implement its methods.

This ensures that the Memory class can interact with any backend
without knowing its internal details.

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any, Iterable, Iterator, Optional
from dataclasses import asdict
import logging
from hieromem_core.episode import Episode


class MemoryStore(ABC):
    """
    Abstract base class defining the persistence contract.

    Each subclass must implement CRUD operations for storing and
    retrieving Episode objects, while preserving the theoretical
    semantics of the HIEROMEM framework.

    Attributes:
        debug: Enable debug logging
        logger: Logger instance for this store
        graph_store: Optional reference to a GraphStore (set by Memory or self for GraphStore subclasses)
    """

    # ------------------------------------------------------------------ #
    # Constructor / Logging utilities
    # ------------------------------------------------------------------ #
    def __init__(self, debug: bool = False):
        self.debug = debug
        self.logger = logging.getLogger(self.__class__.__name__)
        # Graph store attribute for Memory compatibility
        # Subclasses that implement GraphStore should set this to self
        self.graph_store: Optional["MemoryStore"] = None
        if self.debug and not self.logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(
                logging.Formatter("[%(levelname)s] %(name)s: %(message)s")
            )
            self.logger.addHandler(handler)
            self.logger.setLevel(logging.INFO)

    def _log(self, message: str) -> None:
        """Internal helper for lightweight debug logging."""
        if self.debug:
            self.logger.info(message)

    # ------------------------------------------------------------------ #
    # Abstract CRUD interface
    # ------------------------------------------------------------------ #
    @abstractmethod
    def save_episode(self, ep: Episode) -> None:
        """Persist a single episode."""

    @abstractmethod
    def save_many(self, episodes: Iterable[Episode]) -> None:
        """Persist multiple episodes."""

    @abstractmethod
    def load_all(self) -> list[Episode]:
        """Load all episodes currently stored."""

    @abstractmethod
    def load_by_id(self, episode_id: str) -> Optional[Episode]:
        """Return a single episode by ID, or None if not found."""

    @abstractmethod
    def update_trust(self, episode_id: str, trust: float) -> None:
        """Update the trust value of a specific episode."""

    @abstractmethod
    def remove(self, episode_id: str) -> None:
        """Delete a single episode."""

    @abstractmethod
    def remove_many(self, episode_ids: Iterable[str]) -> None:
        """Delete multiple episodes."""

    @abstractmethod
    def clear(self) -> None:
        """Remove all stored episodes."""

    # ------------------------------------------------------------------ #
    # Helper utilities shared across backends
    # ------------------------------------------------------------------ #
    @staticmethod
    def episode_to_dict(ep: Episode) -> dict[str, Any]:
        """
        Convert an Episode into a JSON-serializable dictionary.

        Concrete stores can reuse this helper when persisting episodes.
        """
        d = asdict(ep)

        # NetworkX graphs → simple node/edge representation
        G = ep.G
        d["G"] = {
            "nodes": [{"id": n, **(G.nodes[n] or {})} for n in G.nodes],
            "edges": [{"u": u, "v": v, **(G.edges[u, v] or {})} for u, v in G.edges],
        }

        # NumPy vectors → lists
        d["e"] = ep.e.tolist()

        # datetime → ISO string
        d["tocc"] = ep.tocc.isoformat()
        return d

    @staticmethod
    def validate_episode(ep: Episode) -> None:
        """Basic integrity checks for an Episode before persistence."""
        if not (0.0 <= ep.trust <= 1.0):
            raise ValueError(
                f"Invalid trust value {ep.trust} for episode {ep.id}")
        if ep.e is None or len(ep.e) == 0:
            raise ValueError(f"Episode {ep.id} has empty embedding vector.")
        if ep.G is None:
            raise ValueError(f"Episode {ep.id} missing scene graph G.")

    # ------------------------------------------------------------------ #
    # Convenience methods
    # ------------------------------------------------------------------ #
    def __len__(self) -> int:
        """Return the number of stored episodes (for convenience)."""
        try:
            return len(self.load_all())
        except Exception:
            return 0

    def __iter__(self) -> Iterator[Episode]:
        """Allow iteration over stored episodes."""
        return iter(self.load_all())


class InMemoryStore(MemoryStore):
    """
    Simple in-memory implementation of MemoryStore for testing and development.
    """
    def __init__(self) -> None:
        super().__init__(debug=False)
        self._episodes: dict[str, Episode] = {}

    def save_episode(self, ep: Episode) -> None:
        # Store a clone to simulate persistence isolation
        self._episodes[ep.id] = ep.clone(new_id=False)

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

    def load_all(self) -> list[Episode]:
        return [ep.clone(new_id=False) for ep in self._episodes.values()]

    def load_by_id(self, episode_id: str) -> Optional[Episode]:
        ep = self._episodes.get(episode_id)
        return ep.clone(new_id=False) if ep is not None else None

    def update_trust(self, episode_id: str, trust: float) -> None:
        if episode_id in self._episodes:
            self._episodes[episode_id].trust = float(trust)

    def remove(self, episode_id: str) -> None:
        self._episodes.pop(episode_id, None)

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

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