"""
backend_store.py
----------------
Unified MemoryStore implementation delegating to pluggable Vector and Graph backends.
"""

from __future__ import annotations

import json
import numpy as np
import networkx as nx
from datetime import datetime
from typing import Any, Iterable, List, Optional, Tuple, Dict

from hieromem_core.episode import Episode

from hieromem_core.persistence.graph_store import GraphStore
from hieromem_core.backends.interface import VectorBackend, GraphBackend


class BackendMemoryStore(GraphStore):
    """
    Unified store using pluggable backends.

    This class bridges the high-level Episode persistence contract
    with low-level Vector and Graph backend interfaces.
    """

    def __init__(
        self,
        vector_backend: VectorBackend,
        graph_backend: GraphBackend,
        debug: bool = False,
    ):
        super().__init__(debug=debug)
        self.vector = vector_backend
        self.graph = graph_backend
        self.graph_store = self  # We implement GraphStore interface

    # ------------------------------------------------------------------ #
    # Serialization Helpers
    # ------------------------------------------------------------------ #
    def _episode_to_metadata(self, ep: Episode) -> Dict[str, Any]:
        """Serialize episode to metadata dict for vector store."""
        # We store the FULL episode structure in metadata so we can reconstruct it
        # purely from the vector store if needed (e.g. if graph store is just an index).
        # However, for huge graphs, this might be heavy.
        # But HIEROMEM episodes are usually scoped.

        G = ep.G
        graph_data = {
            "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],
        }

        return {
            "id": ep.id,
            "tocc": ep.tocc.isoformat(),
            "parent": ep.parent,
            "trust": float(ep.trust),
            "meta": json.dumps(ep.meta) if ep.meta else "{}",
            "graph": json.dumps(graph_data),  # Serialized graph
        }

    def _metadata_to_episode(
        self, eid: str, vec: np.ndarray, meta: Dict[str, Any]
    ) -> Episode:
        """Reconstruct Episode from vector metadata."""
        tocc = datetime.fromisoformat(meta.get("tocc", datetime.now().isoformat()))
        parent = meta.get("parent")
        trust = float(meta.get("trust", 1.0))
        ep_meta = json.loads(meta.get("meta", "{}"))

        graph_json = meta.get("graph")
        G = nx.DiGraph()
        if graph_json:
            G_data = json.loads(graph_json)
            for n in G_data.get("nodes", []):
                nid = n.pop("id")
                G.add_node(nid, **n)
            for e in G_data.get("edges", []):
                u, v = e.pop("u"), e.pop("v")
                G.add_edge(u, v, **e)

        return Episode(
            id=eid, e=vec, G=G, tocc=tocc, meta=ep_meta, parent=parent, trust=trust
        )

    # ------------------------------------------------------------------ #
    # CRUD - MemoryStore
    # ------------------------------------------------------------------ #
    def save_episode(self, ep: Episode) -> None:
        self.validate_episode(ep)

        # 1. Save to Vector Backend
        metadata = self._episode_to_metadata(ep)
        self.vector.add([ep.id], np.array([ep.e]), [metadata])

        # 2. Save to Graph Backend
        triples = []
        for u, v, data in ep.G.edges(data=True):
            # Normalize relation
            rel = (
                data.get("relation_type")
                or data.get("relation")
                or data.get("rel")
                or "RELATED_TO"
            )
            rel = (
                str(rel).strip().lower()
            )  # backend interface expects canonical? interface doesn't enforce but good practice

            # Backend expects (head, relation, tail, properties)
            # Embed generic props
            props = {
                **data,
                "episode": ep.id,
                "trust": float(ep.trust),
                "tocc": ep.tocc.isoformat(),
            }
            triples.append((u, rel, v, props))

        self.graph.add_triples(triples)
        self._log(f"Saved episode {ep.id} to backends")

    def save_many(self, episodes: Iterable[Episode]) -> None:
        # Naive implementation loop. Adapters batch internally if implemented,
        # but interfaces define batch operations.
        # VectorBackend.add takes lists, so we should batch vector adds.
        ep_list = list(episodes)
        if not ep_list:
            return

        ids = []
        vecs = []
        metas = []
        all_triples = []

        for ep in ep_list:
            self.validate_episode(ep)
            ids.append(ep.id)
            vecs.append(ep.e)
            metas.append(self._episode_to_metadata(ep))

            for u, v, data in ep.G.edges(data=True):
                rel = data.get("relation_type") or data.get("relation") or "RELATED_TO"
                rel = str(rel).strip().lower()
                props = {
                    **data,
                    "episode": ep.id,
                    "trust": ep.trust,
                    "tocc": ep.tocc.isoformat(),
                }
                all_triples.append((u, rel, v, props))

        self.vector.add(ids, np.array(vecs), metas)
        self.graph.add_triples(all_triples)
        self._log(f"Saved {len(ep_list)} episodes (batched)")

    def load_all(self) -> List[Episode]:
        # VectorBackend needs a scan/iterate capability.
        # Many vector stores (Pinecone) don't support "dump all" easily.
        # However, for consistency, we rely on the vector backend methods.
        # If vector backend doesn't support full scan, this fails.
        # But Faiss and PGVector do.
        # Warning: This operation is expensive and not scalable.
        raise NotImplementedError(
            "load_all is not scalable for pluggable backends. Use search or ID lookup."
        )

    def load_by_id(self, episode_id: str) -> Optional[Episode]:
        # Fetch from vector backend (primary source of truth for full record)
        results = self.vector.get([episode_id])
        if results and results[0]:
            eid, vec, meta = results[0]
            return self._metadata_to_episode(eid, vec, meta)
        return None

    def update_trust(self, episode_id: str, trust: float) -> None:
        # This requires read-modify-write on vector store metadata
        ep = self.load_by_id(episode_id)
        if ep:
            ep.trust = trust
            self.save_episode(ep)  # Re-save with updated trust

    def remove(self, episode_id: str) -> None:
        self.vector.delete([episode_id])
        # Graph deletion is harder: need to remove edges associated with this episode.
        # GraphBackend.delete_triples takes specific triples.
        # To delete efficiently, GraphBackend needs "delete_by_property" or "delete_subgraph".
        # Current interface is minimal.
        # We can implement 'search_triples' to find edges with 'episode' property?
        # GraphBackend.search_triples doesn't support property filtering in current abstract interface signature
        # (it returns properties, but inputs are head/rel/tail).
        # We rely on graph backend specific capabilities or ignore graph cleanup for now (GC later).
        # Let's assume eventual consistency or manual cleanup.
        pass

    def remove_many(self, episode_ids: Iterable[str]) -> None:
        self.vector.delete(list(episode_ids))

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

    # ------------------------------------------------------------------ #
    # Query - GraphStore
    # ------------------------------------------------------------------ #
    def search_similar(
        self, query_vec: np.ndarray, k: int = 5
    ) -> List[Tuple[str, float]]:
        results = self.vector.search(query_vec, k=k)
        return [(r[0], r[1]) for r in results]

    def query_relations(
        self, entity: str, depth: int = 1, *, namespace: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        # Map GraphBackend.get_neighbors -> list[dict]
        # GraphBackend.get_neighbors returns (u, rel, v, props)
        neighbors = self.graph.get_neighbors(entity, depth=depth)
        results = []
        for u, rel, v, props in neighbors:
            results.append(
                {
                    "src": u,
                    "dst": v,
                    "relation": rel,
                    "trust": props.get("trust", 1.0),
                    "properties": props,
                }
            )
        return results

    def search_related(
        self,
        episode_id: str,
        depth: int = 2,
        limit: int = 50,
        *,
        namespace: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        # This requires finding the graph nodes for the episode first.
        # Then resolving neighbors.
        # If episode is not a node itself, we look at its content nodes.
        # Assuming episode ID itself might be a node or connected?
        # Usually we link episode -> observed -> entities.
        # Let's assume we query neighbors of episode ID.
        neighbors = self.graph.get_neighbors(episode_id, depth=depth)
        # Convert to flattened structure
        results = []
        for u, rel, v, props in neighbors:
            results.append(
                {
                    "source": u,
                    "target": v,
                    "relations": [rel],
                    "trust": props.get("trust", 1.0),
                    "properties": props,
                }
            )
        return results[:limit]

    def get_graph_stats(self, *, namespace: Optional[str] = None) -> Dict[str, Any]:
        return {"backend": "pluggable", "info": "stats_not_implemented"}
