"""
# flake8: noqa
recall_service.py
-----------------
High-level orchestration logic for episodic retrieval (the "recall" operation)
in the HIEROMEM architecture.

Bridges:
  • FastAPI layer (user query requests)
  • HIEROMEM Core (Memory, Encoder, Episode)
  • Response formatting (schemas + metrics updates)

Implements trust-weighted, content-addressable retrieval using resonance
mechanics, as formalized in the theoretical framework.

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations

from collections import OrderedDict
from datetime import datetime
from functools import lru_cache
from typing import Any, Dict, List
import hashlib
import logging

import numpy as np

from hieromem_core.episode import Episode
from hieromem_core.memory import Memory
from hieromem_core.services.cache_service import CacheService
from hieromem_api.instrumentation.prometheus import (
    update_episode_metrics,
    update_stability_metrics,
)
from hieromem_api.instrumentation.tracing import trace_function
from hieromem_api.models.responses import EpisodeSummary, RecallResponse


class RecallService:
    """
    Provides a clean API for retrieving semantically relevant memories
    given a natural language query.

    Handles:
      • Embedding the query via encoder
      • Resonance-weighted retrieval
      • Conversion to standardized API schema
      • Metrics updates and tracing
      • Distributed caching via Redis (with in-memory fallback)
    """

    def __init__(
        self,
        memory: Memory,
        tracer: Any | None = None,
        cache_service: CacheService | None = None,
        *,
        cache_ttl_seconds: int | None = None,
    ):
        self.memory = memory
        self.tracer = tracer
        
        # Distributed cache (Redis) with in-memory fallback
        self.cache_service = cache_service or CacheService()
        self._cache_ttl_seconds = cache_ttl_seconds or 3600
        self._cache_versions: dict[str, int] = {}
        
        # Keep in-memory cache as L1 for single-node deployments or when Redis is unavailable
        self._result_cache: OrderedDict[
            tuple[str, int, str | None], tuple[tuple[int, str], RecallResponse]
        ] = OrderedDict()
        self._result_cache_size = 128
        
        # Subscribe to cache invalidation events from other instances
        if self.cache_service:
            self.cache_service.subscribe_to_invalidations(self._handle_remote_invalidation)

    # ----------------------------------------------------------------------- #
    # Main recall entrypoint
    # ----------------------------------------------------------------------- #

    @trace_function(tracer=None, name="recall_query")
    def recall(
            self,
            query: str,
            k: int = 3,
            *,
            namespace: str | None = None) -> RecallResponse:
        """
        Perform a semantic recall query and return top-k episodes.

        Args:
            query: User query or natural-language cue.
            k: Number of most similar episodes to return.

        Returns:
            RecallResponse: structured retrieval result.
        """
        if not hasattr(self.memory, "encoder"):
            raise RuntimeError(
                "Memory encoder not initialized in core engine.")

        cache_key = self._cache_key(query, k, namespace)
        revision = getattr(self.memory, "revision", 0)
        cache_version = self._current_cache_version(namespace)
        version = (revision, cache_version)
        
        # L2 Cache: Check Redis (distributed)
        redis_key = self._redis_cache_key(cache_key, version)
        cached_response = self.cache_service.get(redis_key)

        if isinstance(cached_response, dict):
            # Cache hit in Redis - deserialize and return
            response = RecallResponse(**cached_response)
            self._record_cache_hit(namespace, "L2")
        else:
            # L1 Cache: Check in-memory (single-node fallback)
            cached = self._result_cache.get(cache_key)
            if cached and cached[0] == version:
                response = cached[1].model_copy(deep=True)
                self._record_cache_hit(namespace, "L1")
            else:
                # Cache miss - perform actual retrieval
                response = self._recall_uncached(query, k, namespace)
                self._remember_cached_result(cache_key, version, response)
                self._record_cache_miss(namespace)

        self._update_metrics(namespace)
        return response

    # ----------------------------------------------------------------------- #
    # Explanation utilities
    # ----------------------------------------------------------------------- #

    def explain(self, query: str, k: int = 3, *
    , namespace: str | None = None) -> Dict[str, Any]:
        """
        Return a human-readable explanation of recall mechanics for diagnostics.
        """
        response = self.recall(query, k, namespace=namespace)
        return {
            "query": response.query,
            "episodes_returned": len(response.results),
            "average_trust": sum(e.trust for e in response.results) / max(
                1, len(response.results)
            ),
            "episodes": [e.model_dump() for e in response.results],
        }

    # ----------------------------------------------------------------------- #
    # Internal helpers
    # ----------------------------------------------------------------------- #

    def _recall_uncached(
            self,
            query: str,
            k: int,
            namespace: str | None) -> RecallResponse:
        query_vec = self._encode_cached(query)
        retrieved_eps: List[Episode] = self.memory.resonant_retrieve(
            query_vec, k=k, namespace=namespace)
        if not retrieved_eps:
            if getattr(self.memory, "_multi_tenant", False):
                ns_map = getattr(self.memory, "_episodes_by_namespace", {})
                retrieved_eps = list(ns_map.get(namespace, []))
                if not retrieved_eps:
                    retrieved_eps = list(getattr(self.memory, "episodes", []))
            else:
                retrieved_eps = list(getattr(self.memory, "episodes", []))
            retrieved_eps = retrieved_eps[:k]
        results = [
            EpisodeSummary(
                id=ep.id,
                trust=float(ep.trust),
                timestamp=self._safe_timestamp(ep),
                meta=self._safe_meta(ep),
            )
            for ep in retrieved_eps
        ]
        return RecallResponse(query=query, k=k, results=results)

    def _safe_timestamp(self, ep: Episode) -> str:
        raw_ts = getattr(ep, "tocc", None) or getattr(ep, "t", None) or getattr(ep, "timestamp", None)
        if raw_ts and hasattr(raw_ts, "isoformat"):
            try:
                candidate = raw_ts.isoformat()
                if isinstance(candidate, str):
                    return candidate
            except Exception as exc:  # pragma: no cover - defensive fallback
                logging.getLogger(__name__).warning(
                    "Failed to serialise timestamp %s: %s", raw_ts, exc
                )
        return datetime.utcnow().isoformat()

    def _safe_meta(self, ep: Episode) -> dict:
        meta = getattr(ep, "meta", {}) or {}
        return meta if isinstance(meta, dict) else {}

    def _update_metrics(self, namespace: str | None) -> None:
        try:
            with self.memory.use_namespace(namespace):
                total_eps = len(self.memory.episodes)
                avg_trust = (sum(ep.trust for ep in self.memory.episodes) /
                             total_eps if total_eps > 0 else 0.0)
                snapshot = getattr(self.memory, "stability_monitor", None)
                if snapshot is None:
                    return
                snapshot = snapshot.stability()
        except Exception as exc:  # pragma: no cover - metrics best-effort
            logging.getLogger(__name__).warning(
                "Failed to update metrics for namespace %s: %s", namespace, exc
            )
            return

        ns = namespace or getattr(self.memory, "namespace", "default")
        update_episode_metrics(
            total_eps,
            avg_trust,
            namespace=ns,
            capacity=self.memory.Nmax,
        )
        update_stability_metrics(
            namespace=ns,
            distortion_bound=snapshot.summary_distortion_bound,
            contraction_factor=snapshot.contraction_factor,
            lyapunov_margin=snapshot.lyapunov_margin,
            lyapunov_warning=snapshot.lyapunov_warning,
        )

    def _cache_key(self, query: str, k: int, namespace: str |
                   None) -> tuple[str, int, str | None]:
        return (query, int(k), namespace or None)

    def _remember_cached_result(
        self,
        key: tuple[str, int, str | None],
        version: tuple[int, str],
        response: RecallResponse,
    ) -> None:
        # Write to L2 (Redis) - distributed cache
        redis_key = self._redis_cache_key(key, version)
        self.cache_service.set(
            redis_key,
            response.model_dump(),
            ttl=self._cache_ttl_seconds,
        )
        
        # Write to L1 (in-memory) - local cache
        self._result_cache[key] = (version, response.model_copy(deep=True))
        self._result_cache.move_to_end(key)
        while len(self._result_cache) > self._result_cache_size:
            self._result_cache.popitem(last=False)

    def invalidate_cache(self, namespace: str | None = None, broadcast: bool = True) -> None:
        """Drop cached recall responses, optionally scoped to a namespace.
        
        Args:
            namespace: Namespace to invalidate, or None for all
            broadcast: If True, publish invalidation event to other instances via Redis pub/sub
        """

        if namespace is None:
            # Clear all caches
            self._result_cache.clear()
            # Redis: We need to track keys or use a pattern, for now clear locally
            # In production, consider using Redis key patterns or a separate invalidation channel
            if broadcast and self.cache_service:
                version = self._bump_cache_version("__all__")
                self.cache_service.publish_invalidation("__all__", version=version)
            return

        # Clear namespace-specific entries from in-memory cache
        keys_to_remove = [
            key for key in self._result_cache if key[2] == namespace]
        for key in keys_to_remove:
            self._result_cache.pop(key, None)
            # Also clear from Redis
            # We need to reconstruct all possible version keys - in practice, use current version
            revision = getattr(self.memory, "revision", 0)
            cache_version = self._current_cache_version(namespace)
            redis_key = self._redis_cache_key(key, (revision, cache_version))
            self.cache_service.delete(redis_key)
        
        # Broadcast invalidation to other instances
        if broadcast and self.cache_service:
            version = self._bump_cache_version(namespace)
            self.cache_service.publish_invalidation(namespace, version=version)

    def _redis_cache_key(
        self,
        cache_key: tuple[str, int, str | None],
        version: tuple[int, str],
    ) -> str:
        """Generate a Redis cache key from the tuple cache key and version."""
        query, k, namespace = cache_key
        revision, cache_version = version
        # Create a deterministic hash of the query to keep keys short
        query_hash = hashlib.sha256(query.encode(), usedforsecurity=False).hexdigest()[:16]
        ns_part = f":{namespace}" if namespace else ""
        return f"recall:r{revision}:cv{cache_version}:{query_hash}:k{k}{ns_part}"

    def _record_cache_hit(self, namespace: str | None, cache_tier: str) -> None:
        """Record a cache hit to Prometheus metrics."""
        try:
            from hieromem_api.instrumentation.prometheus import record_cache_hit
            record_cache_hit(namespace, cache_tier)
        except ImportError:
            pass  # Metrics not available

    def _record_cache_miss(self, namespace: str | None) -> None:
        """Record a cache miss to Prometheus metrics."""
        try:
            from hieromem_api.instrumentation.prometheus import record_cache_miss
            record_cache_miss(namespace)
        except ImportError:
            pass  # Metrics not available

    def _handle_remote_invalidation(self, namespace: str, version: int | None) -> None:
        """Handle cache invalidation event from another instance via Redis pub/sub."""
        if version is not None:
            self._cache_versions[namespace] = version
        # Invalidate local cache without broadcasting (to avoid infinite loop)
        if namespace == "__all__":
            self.invalidate_cache(namespace=None, broadcast=False)
        else:
            self.invalidate_cache(namespace=namespace, broadcast=False)

    def _current_cache_version(self, namespace: str | None) -> str:
        global_version = self._get_cache_version("__all__")
        if namespace is None:
            return str(global_version)
        namespace_version = self._get_cache_version(namespace)
        return f"{global_version}.{namespace_version}"

    def _get_cache_version(self, namespace: str) -> int:
        if namespace in self._cache_versions:
            return self._cache_versions[namespace]
        version = self.cache_service.get_version(namespace)
        self._cache_versions[namespace] = version
        return version

    def _bump_cache_version(self, namespace: str) -> int:
        version = self.cache_service.bump_version(namespace)
        self._cache_versions[namespace] = version
        return version

    @lru_cache(maxsize=1024)
    def _encode_cached(self, query: str) -> np.ndarray:
        vector = self.memory.encoder.encode_single(query)
        return np.asarray(vector, dtype=np.float32)
