"""Graph-only search helper built on top of knowledge graph traversal."""

from __future__ import annotations

from collections import deque

from neurovault.cognitive.entity_extractor import EntityExtractor
from neurovault.search.multi_signal import RecallResult
from neurovault.storage import NeurovaultStorage


class GraphSearchEngine:
    """Entity-driven graph traversal memory search."""

    def __init__(self, storage: NeurovaultStorage, bridge, extractor: EntityExtractor | None = None):
        self._storage = storage
        self._bridge = bridge
        self._extractor = extractor or EntityExtractor()

    def search(self, query: str, limit: int = 20, min_importance: float = 0.0) -> list[RecallResult]:
        query_entities = self._extractor.extract(query)
        if not query_entities:
            return []

        memory_scores: dict[bytes, float] = {}
        for query_entity in query_entities:
            matches = self._storage.find_entities_by_name(query_entity["name"])
            for entity in matches:
                for node in self.traverse(entity.id, max_depth=2, max_results=30, min_strength=0.3):
                    for memory_id, confidence in self._storage.get_memories_for_entity(node["entity"].id):
                        depth_penalty = 1.0 / (1 + node["depth"])
                        memory_scores[memory_id] = memory_scores.get(memory_id, 0.0) + confidence * depth_penalty

        results: list[RecallResult] = []
        for memory_id in sorted(memory_scores, key=lambda mid: -memory_scores[mid])[:limit]:
            try:
                memory = self._bridge.get_memory(memory_id)
            except Exception:
                continue
            if memory.importance < min_importance:
                continue
            results.append(
                RecallResult(
                    memory=memory,
                    score=memory_scores[memory_id],
                    match_type="graph",
                    entities=self._storage.get_entities_for_memory(memory_id),
                )
            )
        return results

    def traverse(
        self,
        start_entity_id: bytes,
        max_depth: int = 3,
        max_results: int = 50,
        min_strength: float = 0.3,
    ) -> list[dict]:
        visited = set()
        queue = deque([(start_entity_id, 0)])
        results = []

        while queue and len(results) < max_results:
            entity_id, depth = queue.popleft()
            if entity_id in visited or depth > max_depth:
                continue
            visited.add(entity_id)

            try:
                entity = self._storage.get_entity(entity_id)
            except Exception:
                continue

            results.append({"entity": entity, "depth": depth})

            for rel in self._storage.get_relationships_from(entity_id, min_strength):
                if rel.target_entity_id not in visited:
                    queue.append((rel.target_entity_id, depth + 1))

            for rel in self._storage.get_relationships_to(entity_id, min_strength):
                if rel.source_entity_id not in visited:
                    queue.append((rel.source_entity_id, depth + 1))

        return results
