"""
H-MEM Retrieval: Semantic search across memory hierarchy.

Provides HMEMRetrieval class for searching episodes and traces
via semantic similarity with graceful degradation when no
embedding model is available.
"""

import json
import math
import re
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from gsd_rlm.memory.hmem.store import EpisodeStore


@dataclass
class SearchResult:
    """A search result with memory_id, similarity score, and metadata."""

    memory_id: str
    similarity_score: float
    memory_type: str  # "episode" or "trace"
    content: str  # Summary or context
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "memory_id": self.memory_id,
            "similarity_score": self.similarity_score,
            "memory_type": self.memory_type,
            "content": self.content,
            "metadata": self.metadata,
        }


class HMEMRetrieval:
    """
    Semantic search across H-MEM hierarchy (MEM-05).

    Searches episodes and traces using semantic similarity when
    embeddings are available, falls back to keyword matching otherwise.

    Attributes:
        episode_store: Store for episode retrieval.
        embedding_model: Optional callable that returns embeddings for text.
        db_path: Path to SQLite database for direct queries.
        embedding_cache: Cache for frequently accessed embeddings.

    Example:
        >>> store = EpisodeStore("memory/hmem.db")
        >>> retrieval = HMEMRetrieval(store)
        >>> results = await retrieval.search("code review issues", k=5)
    """

    # Valid levels for search
    VALID_LEVELS = {"episode", "trace", "category", "domain"}
    DEFAULT_LEVELS = ["episode", "trace"]

    def __init__(
        self,
        episode_store: Optional["EpisodeStore"] = None,
        embedding_model: Optional[Any] = None,
        db_path: str = "memory/hmem.db",
    ):
        """
        Initialize HMEMRetrieval.

        Args:
            episode_store: EpisodeStore for retrieving episodes.
            embedding_model: Optional model that provides embeddings.
                            Can be any callable that takes a string and returns
                            List[float]. If None, uses keyword matching.
            db_path: Path to SQLite database for direct queries.
        """
        self.episode_store = episode_store
        self.embedding_model = embedding_model
        self.db_path = Path(db_path)
        self.embedding_cache: Dict[str, List[float]] = {}
        self._cache_max_size = 1000

    async def search(
        self,
        query: str,
        k: int = 5,
        levels: Optional[List[str]] = None,
    ) -> List[Tuple[str, float]]:
        """
        Search for memories similar to query.

        Args:
            query: Search query string.
            k: Maximum number of results to return.
            levels: Memory levels to search ["episode", "trace", "category", "domain"].
                   Default: ["episode", "trace"]

        Returns:
            List of (memory_id, similarity_score) tuples sorted by score descending.
        """
        if levels is None:
            levels = self.DEFAULT_LEVELS

        # Validate levels
        levels = [l.lower() for l in levels]
        invalid = set(levels) - self.VALID_LEVELS
        if invalid:
            raise ValueError(f"Invalid levels: {invalid}. Valid: {self.VALID_LEVELS}")

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

        results: List[Tuple[str, float]] = []

        # Search episodes
        if "episode" in levels:
            episode_results = await self._search_episodes(query_embedding, k)
            results.extend(episode_results)

        # Search traces
        if "trace" in levels:
            trace_results = await self._search_traces(query_embedding, k)
            results.extend(trace_results)

        # Sort by score descending and take top k
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:k]

    async def search_with_content(
        self,
        query: str,
        k: int = 5,
        levels: Optional[List[str]] = None,
    ) -> List[SearchResult]:
        """
        Search and return full results with content.

        Args:
            query: Search query string.
            k: Maximum number of results.
            levels: Memory levels to search.

        Returns:
            List of SearchResult objects with full content.
        """
        basic_results = await self.search(query, k, levels)
        full_results: List[SearchResult] = []

        for memory_id, score in basic_results:
            if memory_id.startswith("ep_"):
                content, metadata = self._get_episode_content(memory_id)
                full_results.append(
                    SearchResult(
                        memory_id=memory_id,
                        similarity_score=score,
                        memory_type="episode",
                        content=content,
                        metadata=metadata,
                    )
                )
            elif memory_id.startswith("trace_"):
                content, metadata = self._get_trace_content(memory_id)
                full_results.append(
                    SearchResult(
                        memory_id=memory_id,
                        similarity_score=score,
                        memory_type="trace",
                        content=content,
                        metadata=metadata,
                    )
                )

        return full_results

    async def _search_episodes(
        self, query_embedding: Optional[List[float]], k: int
    ) -> List[Tuple[str, float]]:
        """
        Search episodes by embedding similarity.

        Args:
            query_embedding: Query embedding vector (or None for keyword search).
            k: Maximum results.

        Returns:
            List of (episode_id, score) tuples.
        """
        if not self.db_path.exists():
            return []

        conn = sqlite3.connect(str(self.db_path))
        try:
            cursor = conn.cursor()

            # Get episodes with embeddings
            cursor.execute(
                """
                SELECT episode_id, embedding, context, action, outcome
                FROM episodes
                WHERE embedding IS NOT NULL
                ORDER BY timestamp DESC
                LIMIT ?
                """,
                (k * 3,),  # Get more than needed for filtering
            )
            rows = cursor.fetchall()

            results: List[Tuple[str, float]] = []

            for episode_id, embedding_blob, context, action, outcome in rows:
                if embedding_blob is None:
                    continue

                # Parse embedding
                try:
                    episode_embedding = json.loads(embedding_blob.decode("utf-8"))
                except (json.JSONDecodeError, UnicodeDecodeError):
                    continue

                # Calculate similarity
                if query_embedding is not None:
                    score = self._cosine_similarity(query_embedding, episode_embedding)
                else:
                    # Keyword fallback: score based on content match
                    score = self._keyword_score(
                        [context, action, outcome], str(query_embedding)
                    )

                if score > 0:
                    results.append((episode_id, score))

            # Sort and return top k
            results.sort(key=lambda x: x[1], reverse=True)
            return results[:k]
        finally:
            conn.close()

    async def _search_traces(
        self, query_embedding: Optional[List[float]], k: int
    ) -> List[Tuple[str, float]]:
        """
        Search traces by summary embedding similarity.

        Args:
            query_embedding: Query embedding vector (or None for keyword search).
            k: Maximum results.

        Returns:
            List of (trace_id, score) tuples.
        """
        # Check if traces table exists
        if not self.db_path.exists():
            return []

        conn = sqlite3.connect(str(self.db_path))
        try:
            cursor = conn.cursor()

            # Check if traces table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='traces'"
            )
            if not cursor.fetchone():
                return []

            # Get traces with embeddings (if stored)
            cursor.execute(
                """
                SELECT trace_id, summary, theme, patterns_identified, lessons_learned
                FROM traces
                ORDER BY created_at DESC
                LIMIT ?
                """,
                (k * 3,),
            )
            rows = cursor.fetchall()

            results: List[Tuple[str, float]] = []

            for trace_id, summary, theme, patterns, lessons in rows:
                # Combine text for keyword matching
                text_parts = [summary or "", theme or ""]
                if patterns:
                    try:
                        text_parts.extend(json.loads(patterns))
                    except json.JSONDecodeError:
                        pass
                if lessons:
                    try:
                        text_parts.extend(json.loads(lessons))
                    except json.JSONDecodeError:
                        pass

                # For traces without embeddings, use keyword scoring
                if query_embedding is not None:
                    # Use summary for semantic matching if embedding available
                    # For now, use keyword scoring (traces typically don't have embeddings)
                    score = self._keyword_score(text_parts, "")
                else:
                    score = self._keyword_score(text_parts, "")

                if score > 0:
                    results.append((trace_id, score))

            # Sort and return top k
            results.sort(key=lambda x: x[1], reverse=True)
            return results[:k]
        finally:
            conn.close()

    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

        # Check cache
        cache_key = text[:100]  # Use first 100 chars as key
        if cache_key in self.embedding_cache:
            return self.embedding_cache[cache_key]

        # Generate embedding
        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 []
            else:
                return None

            # Cache result
            if len(self.embedding_cache) < self._cache_max_size:
                self.embedding_cache[cache_key] = embedding

            return embedding
        except Exception:
            return None

    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 score (0.0 to 1.0).
        """
        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

        return dot_product / (norm_a * norm_b)

    def _keyword_score(self, text_parts: List[str], query: str) -> float:
        """
        Calculate keyword-based relevance score.

        Args:
            text_parts: List of text content to search.
            query: Search query (ignored for now, uses all terms).

        Returns:
            Relevance score based on content richness.
        """
        combined = " ".join(text_parts).lower()
        if not combined.strip():
            return 0.0

        # Score based on content length and diversity
        words = set(re.findall(r"\w+", combined))
        if not words:
            return 0.0

        # Base score from content richness
        score = min(1.0, len(words) / 50)  # Max at 50 unique words

        # Boost for action verbs
        action_verbs = {
            "created",
            "implemented",
            "fixed",
            "added",
            "removed",
            "updated",
            "refactored",
            "tested",
            "analyzed",
            "resolved",
        }
        if words & action_verbs:
            score += 0.2

        return min(1.0, score)

    def _get_episode_content(self, episode_id: str) -> Tuple[str, Dict[str, Any]]:
        """
        Get episode content and metadata by ID.

        Args:
            episode_id: Episode identifier.

        Returns:
            Tuple of (content string, metadata dict).
        """
        if self.episode_store:
            episode = self.episode_store.get(episode_id)
            if episode:
                content = f"{episode.context}\n{episode.action}\n{episode.outcome}"
                return content, {
                    "agent_id": episode.agent_id,
                    "session_id": episode.session_id,
                    "episode_type": episode.episode_type.value,
                    "success": episode.success,
                    "tags": episode.tags,
                }

        # Fallback to direct DB query
        if self.db_path.exists():
            conn = sqlite3.connect(str(self.db_path))
            try:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    SELECT context, action, outcome, agent_id, session_id, episode_type, success, tags
                    FROM episodes WHERE episode_id = ?
                    """,
                    (episode_id,),
                )
                row = cursor.fetchone()
                if row:
                    (
                        context,
                        action,
                        outcome,
                        agent_id,
                        session_id,
                        ep_type,
                        success,
                        tags,
                    ) = row
                    content = f"{context or ''}\n{action or ''}\n{outcome or ''}"
                    try:
                        tags_list = json.loads(tags) if tags else []
                    except json.JSONDecodeError:
                        tags_list = []
                    return content, {
                        "agent_id": agent_id,
                        "session_id": session_id,
                        "episode_type": ep_type,
                        "success": bool(success),
                        "tags": tags_list,
                    }
            finally:
                conn.close()

        return "", {}

    def _get_trace_content(self, trace_id: str) -> Tuple[str, Dict[str, Any]]:
        """
        Get trace content and metadata by ID.

        Args:
            trace_id: Trace identifier.

        Returns:
            Tuple of (content string, metadata dict).
        """
        if not self.db_path.exists():
            return "", {}

        conn = sqlite3.connect(str(self.db_path))
        try:
            cursor = conn.cursor()

            # Check if traces table exists
            cursor.execute(
                "SELECT name FROM sqlite_master WHERE type='table' AND name='traces'"
            )
            if not cursor.fetchone():
                return "", {}

            cursor.execute(
                """
                SELECT theme, summary, patterns_identified, lessons_learned,
                       agent_id, session_id, episode_count, overall_success
                FROM traces WHERE trace_id = ?
                """,
                (trace_id,),
            )
            row = cursor.fetchone()
            if row:
                (
                    theme,
                    summary,
                    patterns,
                    lessons,
                    agent_id,
                    session_id,
                    ep_count,
                    success,
                ) = row
                content = f"{theme or ''}\n{summary or ''}"
                try:
                    patterns_list = json.loads(patterns) if patterns else []
                except json.JSONDecodeError:
                    patterns_list = []
                try:
                    lessons_list = json.loads(lessons) if lessons else []
                except json.JSONDecodeError:
                    lessons_list = []
                return content, {
                    "agent_id": agent_id,
                    "session_id": session_id,
                    "episode_count": ep_count,
                    "overall_success": bool(success),
                    "patterns": patterns_list,
                    "lessons": lessons_list,
                }
        finally:
            conn.close()

        return "", {}

    def clear_cache(self) -> None:
        """Clear the embedding cache."""
        self.embedding_cache.clear()
