from __future__ import annotations
from typing import Iterable, Optional

import networkx as nx
import numpy as np

import pytest

from hieromem_core.memory import Memory
from hieromem_core.episode import Episode
from hieromem_core.persistence.memory_store import MemoryStore


def expect(condition: bool, message: str | None = None) -> None:
    """Fail the test when *condition* is false without using ``assert``."""

    if not condition:
        pytest.fail(message or "Expectation failed", pytrace=True)


class FakeIndex:
    def __init__(self) -> None:
        self.vectors: dict[str, np.ndarray] = {}
        self.add_calls: list[str] = []
        self.remove_calls: list[str] = []
        self.search_queries: list[np.ndarray] = []

    @property
    def ids(self) -> list[str]:
        return list(self.vectors.keys())

    def add(self, episode_id: str, vec: np.ndarray) -> None:
        self.vectors[episode_id] = np.asarray(vec, dtype=np.float32)
        self.add_calls.append(episode_id)

    def add_embeddings(self, ids: Iterable[str], matrix: np.ndarray) -> None:
        """Support the Memory helper that prefers add_embeddings for batches."""

        for episode_id, row in zip(ids, matrix):
            self.add(episode_id, row)

    def add_many(self, ids: Iterable[str], matrix: np.ndarray) -> None:
        for episode_id, row in zip(ids, matrix):
            self.add(episode_id, row)

    def remove(self, episode_id: str) -> None:
        self.vectors.pop(episode_id, None)
        self.remove_calls.append(episode_id)

    def search(self, query: np.ndarray, k: int) -> list[tuple[str, float]]:
        self.search_queries.append(np.asarray(query, dtype=np.float32))
        if not self.vectors:
            return []
        scored = [(eid, float(np.dot(vec, query)))
                  for eid, vec in self.vectors.items()]
        scored.sort(key=lambda pair: pair[1], reverse=True)
        return scored[:k]

    def clear(self) -> None:
        self.vectors.clear()
        self.add_calls.clear()
        self.remove_calls.clear()
        self.search_queries.clear()


class InMemoryStore(MemoryStore):
    def __init__(self) -> None:
        super().__init__(debug=False)
        self._episodes: dict[str, Episode] = {}

    def save_episode(self, ep: Episode) -> None:
        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()


def test_insert_updates_index_with_episode_vector() -> None:
    memory = Memory(Nmax=5)
    fake_index = FakeIndex()
    memory._index = fake_index

    embedding = np.array([0.5, 0.5, 0.5], dtype=np.float32)
    memory.insert(embedding, nx.DiGraph(), meta={"label": "demo"})

    expect(bool(memory.episodes), "insert should create an episode")
    episode_id = memory.episodes[0].id
    expect(episode_id in fake_index.vectors,
           "episode vector should be stored in index")
    expect(fake_index.add_calls == [episode_id],
           "index add should be called once with episode id")


def test_retrieve_uses_index_and_drops_stale_ids() -> None:
    store = InMemoryStore()
    memory = Memory(Nmax=5, store=store, tau=0.3)
    fake_index = FakeIndex()
    memory._index = fake_index

    emb_a = np.array([1.0, 0.0, 0.0], dtype=np.float32)
    emb_b = np.array([0.0, 1.0, 0.0], dtype=np.float32)

    ep_a = memory.insert(emb_a, nx.DiGraph())
    memory.insert(emb_b, nx.DiGraph())

    # Inject a stale entry that no longer exists in persistence
    fake_index.add("stale-id", np.array([1.0, 1.0, 1.0], dtype=np.float32))
    store.remove("stale-id")

    results = memory.resonant_retrieve(np.array([0.9, 0.1, 0.0], dtype=np.float32), k=2)

    removed_stale = "stale-id" in fake_index.remove_calls if fake_index.remove_calls else False
    expect(removed_stale, "stale ids should be removed from index")
    expect(bool(results and results[0].id == ep_a.id),
           "retrieval should return closest episode")
    expect(bool(fake_index.search_queries), "index search should be invoked")
    last_event = memory.get_recent_events(1)[0]
    expect(last_event.details.get("backend") == "faiss",
           "event backend should be recorded")


def test_update_trust_removes_evicted_episode_from_index() -> None:
    memory = Memory(Nmax=5, theta=0.5, lam=0.9)
    fake_index = FakeIndex()
    memory._index = fake_index

    ep = memory.insert(np.array([0.1, 0.2], dtype=np.float32), nx.DiGraph())
    memory.episodes[0].trust = 0.1

    memory.update_trust()

    expect(ep.id not in fake_index.ids,
           "evicted episode should be removed from index")
    expect(all(e.id != ep.id for e in memory.episodes),
           "episode list should no longer contain evicted episode")
