"""
Memory Router for GSD-RLM Memory System.

Routes context queries to appropriate memory stores based on semantic
similarity between the query and store descriptions.

Provides MemoryStoreType enum for all available memory stores and
MemoryRouter for intelligent query routing.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional


class MemoryStoreType(Enum):
    """
    Types of memory stores available in the GSD-RLM system.

    Each store type has a specific scope and purpose for memory storage.
    """

    # Session-level stores
    SESSION = "SESSION"  # Current conversation and recent task outputs

    # H-MEM hierarchy
    HMEM_EPISODE = "HMEM_EPISODE"  # Recent agent experiences and actions
    HMEM_TRACE = "HMEM_TRACE"  # Consolidated work sessions and patterns
    HMEM_CATEGORY = "HMEM_CATEGORY"  # Grouped experiences by domain type
    HMEM_DOMAIN = "HMEM_DOMAIN"  # Generalized knowledge and principles

    # Memory Bridge levels
    BRIDGE_L0 = "BRIDGE_L0"  # Current session facts and decisions
    BRIDGE_L1 = "BRIDGE_L1"  # Current phase progress and context
    BRIDGE_L2 = "BRIDGE_L2"  # Project-level facts and requirements
    BRIDGE_L3 = "BRIDGE_L3"  # Cross-project patterns and preferences

    # InfiniRetri
    INFINIRETRI = "INFINIRETRI"  # Large documents and codebases

    @property
    def description(self) -> str:
        """Get human-readable description of this store type."""
        descriptions = {
            MemoryStoreType.SESSION: "current conversation and recent task outputs",
            MemoryStoreType.HMEM_EPISODE: "recent agent experiences and actions",
            MemoryStoreType.HMEM_TRACE: "consolidated work sessions and patterns",
            MemoryStoreType.HMEM_CATEGORY: "grouped experiences by domain type",
            MemoryStoreType.HMEM_DOMAIN: "generalized knowledge and principles",
            MemoryStoreType.BRIDGE_L0: "current session facts and decisions",
            MemoryStoreType.BRIDGE_L1: "current phase progress and context",
            MemoryStoreType.BRIDGE_L2: "project-level facts and requirements",
            MemoryStoreType.BRIDGE_L3: "cross-project patterns and preferences",
            MemoryStoreType.INFINIRETRI: "large documents and codebases",
        }
        return descriptions.get(self, "unknown memory store")

    @property
    def category(self) -> str:
        """Get the category of this store type."""
        if self in (MemoryStoreType.SESSION,):
            return "session"
        elif self in (
            MemoryStoreType.HMEM_EPISODE,
            MemoryStoreType.HMEM_TRACE,
            MemoryStoreType.HMEM_CATEGORY,
            MemoryStoreType.HMEM_DOMAIN,
        ):
            return "hmem"
        elif self in (
            MemoryStoreType.BRIDGE_L0,
            MemoryStoreType.BRIDGE_L1,
            MemoryStoreType.BRIDGE_L2,
            MemoryStoreType.BRIDGE_L3,
        ):
            return "bridge"
        else:
            return "retrieval"


@dataclass
class MemoryRoute:
    """
    A routing decision for a memory query.

    Represents a single store that should be queried, along with
    relevance score and the query embedding used.
    """

    store_type: MemoryStoreType
    relevance_score: float
    query_embedding: Optional[List[float]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "store_type": self.store_type.value,
            "relevance_score": self.relevance_score,
            "query_embedding": self.query_embedding,
        }


class MemoryRouter:
    """
    Routes memory queries to appropriate stores (MEM-08).

    Analyzes query intent and routes to the most relevant memory stores
    based on semantic similarity between query and store descriptions.

    Attributes:
        embedding_model: Optional model for computing embeddings.
        store_embeddings: Pre-computed embeddings for store descriptions.

    Example:
        >>> router = MemoryRouter()
        >>> routes = await router.route("fix the bug", available_stores)
        >>> for route in routes:
        ...     print(f"{route.store_type}: {route.relevance_score}")
    """

    # Store descriptions for relevance matching
    STORE_DESCRIPTIONS = {
        MemoryStoreType.SESSION: "current conversation and recent task outputs",
        MemoryStoreType.HMEM_EPISODE: "recent agent experiences and actions",
        MemoryStoreType.HMEM_TRACE: "consolidated work sessions and patterns",
        MemoryStoreType.HMEM_CATEGORY: "grouped experiences by domain type",
        MemoryStoreType.HMEM_DOMAIN: "generalized knowledge and principles",
        MemoryStoreType.BRIDGE_L0: "current session facts and decisions",
        MemoryStoreType.BRIDGE_L1: "current phase progress and context",
        MemoryStoreType.BRIDGE_L2: "project-level facts and requirements",
        MemoryStoreType.BRIDGE_L3: "cross-project patterns and preferences",
        MemoryStoreType.INFINIRETRI: "large documents and codebases",
    }

    def __init__(self, embedding_model: Optional[Any] = None):
        """
        Initialize the MemoryRouter.

        Args:
            embedding_model: Optional callable that returns embeddings for text.
                            If None, uses keyword matching on descriptions.
        """
        self.embedding_model = embedding_model
        self._store_embeddings: Dict[MemoryStoreType, List[float]] = {}
        self._initialized = False

    async def _ensure_initialized(self) -> None:
        """Initialize store embeddings if embedding model is available."""
        if self._initialized:
            return

        if self.embedding_model is not None:
            for store_type, description in self.STORE_DESCRIPTIONS.items():
                embedding = await self._embed(description)
                if embedding is not None:
                    self._store_embeddings[store_type] = embedding

        self._initialized = True

    async def route(
        self,
        query: str,
        available_stores: Optional[List[MemoryStoreType]] = None,
    ) -> List[MemoryRoute]:
        """
        Route a query to appropriate memory stores.

        Args:
            query: The search query string.
            available_stores: List of stores to consider. If None, uses all.

        Returns:
            List of MemoryRoute objects sorted by relevance_score descending.
        """
        await self._ensure_initialized()

        if available_stores is None:
            available_stores = list(MemoryStoreType)

        # Get query embedding
        query_embedding = await self._embed(query)

        # Calculate relevance for each store
        routes: List[MemoryRoute] = []
        for store_type in available_stores:
            relevance = await self._compute_relevance(
                query_embedding, query, store_type
            )
            if relevance > 0:
                routes.append(
                    MemoryRoute(
                        store_type=store_type,
                        relevance_score=relevance,
                        query_embedding=query_embedding,
                    )
                )

        # Sort by relevance descending
        routes.sort(key=lambda r: r.relevance_score, reverse=True)
        return routes

    async def _compute_relevance(
        self,
        query_embedding: Optional[List[float]],
        query: str,
        store_type: MemoryStoreType,
    ) -> float:
        """
        Compute relevance score for a store against the query.

        Args:
            query_embedding: Pre-computed query embedding (or None).
            query: Original query string.
            store_type: Store type to compute relevance for.

        Returns:
            Relevance score (0.0 to 1.0).
        """
        description = self.STORE_DESCRIPTIONS.get(store_type, "")

        if query_embedding is not None and store_type in self._store_embeddings:
            # Use semantic similarity
            store_embedding = self._store_embeddings[store_type]
            return self._cosine_similarity(query_embedding, store_embedding)
        else:
            # Fall back to keyword matching
            return self._keyword_relevance(query, description, store_type)

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """
        Calculate cosine similarity between two vectors.

        Args:
            a: First vector.
            b: Second vector.

        Returns:
            Cosine similarity (0.0 to 1.0).
        """
        import math

        if not a or not b or len(a) != len(b):
            return 0.0

        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))

        if norm_a == 0 or norm_b == 0:
            return 0.0

        # Normalize to 0-1 range (cosine can be negative)
        similarity = dot_product / (norm_a * norm_b)
        return (similarity + 1) / 2

    def _keyword_relevance(
        self, query: str, description: str, store_type: MemoryStoreType
    ) -> float:
        """
        Calculate keyword-based relevance score.

        Args:
            query: Search query.
            description: Store description.
            store_type: Store type for context-aware scoring.

        Returns:
            Relevance score (0.0 to 1.0).
        """
        query_lower = query.lower()
        desc_lower = description.lower()

        query_words = set(query_lower.split())
        desc_words = set(desc_lower.split())

        # Baseline score for all stores (ensures at least some relevance)
        score = 0.1

        if not query_words:
            return score

        # Base score: word overlap
        overlap = query_words & desc_words
        score += len(overlap) / max(len(query_words), 1)

        # Boost based on specific keywords
        boost = 0.0

        # Session-related keywords
        if store_type == MemoryStoreType.SESSION:
            session_keywords = {"current", "recent", "now", "conversation", "chat"}
            if query_words & session_keywords:
                boost += 0.3

        # Episode-related keywords
        elif store_type == MemoryStoreType.HMEM_EPISODE:
            episode_keywords = {"experience", "action", "happened", "did", "recent"}
            if query_words & episode_keywords:
                boost += 0.3

        # Trace-related keywords
        elif store_type == MemoryStoreType.HMEM_TRACE:
            trace_keywords = {"pattern", "session", "work", "consolidated", "history"}
            if query_words & trace_keywords:
                boost += 0.3

        # Category-related keywords
        elif store_type == MemoryStoreType.HMEM_CATEGORY:
            category_keywords = {"category", "type", "group", "domain", "class"}
            if query_words & category_keywords:
                boost += 0.3

        # Domain-related keywords
        elif store_type == MemoryStoreType.HMEM_DOMAIN:
            domain_keywords = {"knowledge", "principle", "general", "learned", "rule"}
            if query_words & domain_keywords:
                boost += 0.3

        # Bridge L0 (session facts)
        elif store_type == MemoryStoreType.BRIDGE_L0:
            l0_keywords = {"fact", "decision", "current", "session", "now"}
            if query_words & l0_keywords:
                boost += 0.3

        # Bridge L1 (phase context)
        elif store_type == MemoryStoreType.BRIDGE_L1:
            l1_keywords = {"phase", "progress", "context", "current", "stage"}
            if query_words & l1_keywords:
                boost += 0.3

        # Bridge L2 (project facts)
        elif store_type == MemoryStoreType.BRIDGE_L2:
            l2_keywords = {"project", "requirement", "fact", "config", "setting"}
            if query_words & l2_keywords:
                boost += 0.3

        # Bridge L3 (workspace preferences)
        elif store_type == MemoryStoreType.BRIDGE_L3:
            l3_keywords = {"preference", "global", "workspace", "cross-project", "user"}
            if query_words & l3_keywords:
                boost += 0.3

        # InfiniRetri (documents)
        elif store_type == MemoryStoreType.INFINIRETRI:
            infini_keywords = {
                "document",
                "codebase",
                "file",
                "code",
                "large",
                "search",
                "find",
            }
            if query_words & infini_keywords:
                boost += 0.3

        return min(1.0, score + boost)

    async def _embed(self, text: str) -> Optional[List[float]]:
        """
        Generate embedding for text.

        Args:
            text: Text to embed.

        Returns:
            Embedding vector or None if no embedding model.
        """
        if self.embedding_model is None:
            return None

        try:
            if callable(self.embedding_model):
                if hasattr(self.embedding_model, "embed"):
                    embedding = await self.embedding_model.embed(text)
                elif hasattr(self.embedding_model, "encode"):
                    result = self.embedding_model.encode(text)
                    embedding = list(result) if hasattr(result, "__iter__") else []
                else:
                    result = self.embedding_model(text)
                    embedding = list(result) if hasattr(result, "__iter__") else []
                return embedding
        except Exception:
            pass

        return None

    def get_top_stores(
        self, routes: List[MemoryRoute], k: int = 3
    ) -> List[MemoryStoreType]:
        """
        Get top-k store types from routes.

        Args:
            routes: List of MemoryRoute objects.
            k: Number of top stores to return.

        Returns:
            List of MemoryStoreType for top-k stores.
        """
        return [route.store_type for route in routes[:k]]
