"""5-signal search: keyword + semantic + graph + ontology + temporal."""

from __future__ import annotations

import time
from collections import deque
from dataclasses import dataclass, field
from typing import Any

from neurovault.bridge import Memory
from neurovault.cognitive.entity_extractor import EntityExtractor
from neurovault.cognitive.ontology import OntologyEngine
from neurovault.storage import NeurovaultStorage


@dataclass
class RecallResult:
    """Normalized recall result returned by NEUROVAULT search engines."""

    memory: Memory
    score: float
    match_type: str = "multi-signal"
    entities: list[Any] = field(default_factory=list)
    graph_depth: int | None = None


class MultiSignalSearchEngine:
    """5-signal fusion: keyword + semantic + graph + ontology + temporal."""

    RRF_K = 60

    def __init__(self, storage, chronos_bridge, embedding_engine=None, extractor=None, hnsw_index=None):
        self._storage: NeurovaultStorage = storage
        self._bridge = chronos_bridge
        self._embedding = embedding_engine
        self._extractor = extractor or EntityExtractor()
        self._ontology = OntologyEngine(self._storage)
        self._hnsw = hnsw_index

    def attach_hnsw(self, hnsw_index) -> None:
        self._hnsw = hnsw_index

    def search(
        self,
        query: str,
        persona_id: bytes,
        limit: int = 10,
        mode: str = "multi",
        weights: dict | None = None,
        min_importance: float = 0.0,
    ) -> list[RecallResult]:
        if mode == "keyword":
            return self._keyword_only(query, persona_id, limit, min_importance)
        if mode == "semantic":
            return self._semantic_only(query, persona_id, limit, min_importance)
        if mode == "graph":
            return self._graph_only(query, limit, min_importance)
        return self._multi_signal(query, persona_id, limit, weights, min_importance)

    def _multi_signal(
        self,
        query: str,
        persona_id: bytes,
        limit: int,
        weights: dict | None,
        min_importance: float,
    ) -> list[RecallResult]:
        w = weights or {
            "keyword": 1.0,
            "semantic": 1.5,
            "graph": 1.2,
            "ontology": 1.0,
            "temporal": 0.8,
        }

        kw_results = self._keyword_only(query, persona_id, limit * 3, min_importance)
        sem_results = self._semantic_only(query, persona_id, limit * 3, min_importance)
        graph_results = self._graph_only(query, limit * 3, min_importance)
        ontology_results = self._ontology_only(query, persona_id, limit * 3, min_importance)

        scores: dict[bytes, float] = {}
        cache: dict[bytes, Memory] = {}

        for rank, result in enumerate(kw_results):
            memory_id = result.memory.id
            scores[memory_id] = scores.get(memory_id, 0.0) + w["keyword"] / (self.RRF_K + rank + 1)
            cache[memory_id] = result.memory

        for rank, result in enumerate(sem_results):
            memory_id = result.memory.id
            scores[memory_id] = scores.get(memory_id, 0.0) + w["semantic"] / (self.RRF_K + rank + 1)
            cache[memory_id] = result.memory

        for rank, result in enumerate(graph_results):
            memory_id = result.memory.id
            scores[memory_id] = scores.get(memory_id, 0.0) + w["graph"] / (self.RRF_K + rank + 1)
            cache[memory_id] = result.memory

        for rank, result in enumerate(ontology_results):
            memory_id = result.memory.id
            scores[memory_id] = scores.get(memory_id, 0.0) + w["ontology"] / (self.RRF_K + rank + 1)
            cache[memory_id] = result.memory

        now = time.time() * 1000
        for memory_id in scores:
            memory = cache.get(memory_id)
            if memory is None or not hasattr(memory, "created_at"):
                continue
            temporal_weight = self._temporal_weight(memory.created_at, now)
            scores[memory_id] += w["temporal"] * temporal_weight * 0.01

        sorted_ids = sorted(scores.keys(), key=lambda mid: -scores[mid])
        results: list[RecallResult] = []
        for memory_id in sorted_ids[:limit]:
            memory = cache[memory_id]
            entities = self._storage.get_entities_for_memory(memory_id)
            results.append(
                RecallResult(
                    memory=memory,
                    score=scores[memory_id],
                    match_type="multi-signal",
                    entities=entities,
                )
            )
        return results

    def _keyword_only(self, query: str, persona_id: bytes, limit: int, min_importance: float) -> list[RecallResult]:
        del persona_id
        raw_results = self._bridge.search_memories(query, limit, "keyword")
        return [
            RecallResult(memory=r.memory, score=r.score, match_type="keyword")
            for r in raw_results
            if getattr(r.memory, "importance", 0.0) >= min_importance
        ]

    def _semantic_only(self, query: str, persona_id: bytes, limit: int, min_importance: float) -> list[RecallResult]:
        del persona_id
        if self._hnsw is not None and getattr(self._hnsw, "size", 0) > 0:
            query_embedding = self._embed_query(query)
            if query_embedding is not None:
                results: list[RecallResult] = []
                for memory_id, similarity in self._hnsw.search(query_embedding, k=max(1, limit * 2)):
                    try:
                        memory = self._bridge.get_memory(memory_id)
                    except Exception:
                        continue
                    if getattr(memory, "importance", 0.0) < min_importance:
                        continue
                    results.append(RecallResult(memory=memory, score=float(similarity), match_type="semantic-hnsw"))
                    if len(results) >= limit:
                        break
                if results:
                    return results

        raw_results = self._bridge.search_memories(query, limit, "semantic")
        return [
            RecallResult(memory=r.memory, score=r.score, match_type="semantic")
            for r in raw_results
            if getattr(r.memory, "importance", 0.0) >= min_importance
        ]

    def _embed_query(self, query: str) -> list[float] | None:
        if self._embedding is None:
            return None
        try:
            if hasattr(self._embedding, "embed"):
                vector = self._embedding.embed(query)
            elif hasattr(self._embedding, "encode"):
                vector = self._embedding.encode(query)
            else:
                return None
        except Exception:
            return None

        if vector is None:
            return None
        if hasattr(vector, "tolist"):
            vector = vector.tolist()
        if not isinstance(vector, list):
            return None
        return [float(v) for v in vector]

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

        memory_scores: dict[bytes, float] = {}
        graph_depths: dict[bytes, int] = {}

        for query_entity in query_entities:
            matches = self._storage.find_entities_by_name(query_entity["name"])
            for entity in matches:
                traversal = self._traverse(entity.id, max_depth=2, max_nodes=30)
                for item in traversal:
                    linked = self._storage.get_memories_for_entity(item["entity"].id)
                    depth_penalty = 1.0 / (1 + item["depth"])
                    for memory_id, confidence in linked:
                        score = confidence * depth_penalty
                        memory_scores[memory_id] = memory_scores.get(memory_id, 0.0) + score
                        graph_depths[memory_id] = min(graph_depths.get(memory_id, 99), item["depth"])

        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 getattr(memory, "importance", 0.0) < 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),
                    graph_depth=graph_depths.get(memory_id),
                )
            )
        return results

    def _ontology_expand_terms(self, query: str) -> list[str]:
        """Expand query entities with ontology descendants and return new terms only."""
        query_entities = self._extractor.extract(query)
        entity_names = [str(e["name"]) for e in query_entities if e.get("name")]

        if not entity_names:
            for token in query.split():
                token = token.strip().strip(",.!?;:\"'()[]{}")
                if len(token) < 3:
                    continue
                if self._storage.find_entities_by_name(token):
                    entity_names.append(token)

        if not entity_names:
            return []

        expanded = self._ontology.expand_query_with_ontology(entity_names)
        query_lower = query.lower()
        return [name for name in expanded if name.lower() not in query_lower]

    def _ontology_expand_query(self, query: str) -> str:
        """Expand query text with ontology descendants."""
        new_terms = self._ontology_expand_terms(query)
        if not new_terms:
            return query
        return f"{query} {' '.join(new_terms[:5])}"

    def _ontology_only(self, query: str, persona_id: bytes, limit: int, min_importance: float) -> list[RecallResult]:
        """Search using ontology-expanded query terms."""
        new_terms = self._ontology_expand_terms(query)
        if not new_terms:
            return []

        scores: dict[bytes, float] = {}
        cache: dict[bytes, Any] = {}
        for term in new_terms[:5]:
            term_results = self._keyword_only(term, persona_id, limit, min_importance)
            for rank, item in enumerate(term_results):
                memory_id = item.memory.id
                scores[memory_id] = scores.get(memory_id, 0.0) + (1.0 / (self.RRF_K + rank + 1))
                cache[memory_id] = item.memory

        sorted_ids = sorted(scores.keys(), key=lambda mid: -scores[mid])[:limit]
        return [
            RecallResult(
                memory=cache[mid],
                score=scores[mid],
                match_type="ontology",
                entities=self._storage.get_entities_for_memory(mid),
            )
            for mid in sorted_ids
        ]

    def _traverse(self, start_entity_id: bytes, max_depth: int = 2, max_nodes: int = 30) -> list[dict]:
        visited = set()
        queue = deque([(start_entity_id, 0)])
        nodes = []

        while queue and len(nodes) < max_nodes:
            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

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

            for rel in self._storage.get_relationships_from(entity_id, 0.3):
                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, 0.3):
                if rel.source_entity_id not in visited:
                    queue.append((rel.source_entity_id, depth + 1))

        return nodes

    @staticmethod
    def _temporal_weight(created_at: int, now: float) -> float:
        hours = (now - created_at) / (3600 * 1000)
        if hours < 1:
            return 1.0
        if hours < 24:
            return 0.8
        if hours < 168:
            return 0.6
        if hours < 720:
            return 0.4
        return 0.2
