"""
In-memory context-aware Genie cache implementation.

This module provides a context-aware cache that stores embeddings and cache entries
entirely in memory, without requiring external database dependencies like PostgreSQL
or Databricks Lakebase. It uses cosine similarity for similarity search and supports
dual embedding matching (question + conversation context).

Use this when:
- No external database access is available
- Single-instance deployments (cache not shared across instances)
- Cache persistence across restarts is not required
- Cache sizes are moderate (hundreds to low thousands of entries)

For multi-instance deployments or large cache sizes, use PostgresContextAwareGenieService
with PostgreSQL backend instead.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta
from threading import Lock
from typing import Any

import mlflow
import numpy as np
from databricks.sdk import WorkspaceClient
from loguru import logger

from dao_ai.config import (
    GenieInMemoryContextAwareCacheParametersModel,
    WarehouseModel,
)
from dao_ai.genie.cache.base import (
    GenieServiceBase,
    SQLCacheEntry,
)
from dao_ai.genie.cache.context_aware.base import ContextAwareGenieService
from dao_ai.genie.core import GenieResponse


@dataclass
class InMemoryCacheEntry:
    """
    In-memory cache entry storing embeddings and SQL query metadata.

    This dataclass represents a single cache entry stored in memory, including
    pre-normalized dual embeddings (question + context) for fast cosine similarity
    via dot product.

    Uses LRU (Least Recently Used) eviction strategy when capacity is reached.

    Attributes:
        genie_space_id: The Genie space ID this entry belongs to
        question: The original question text
        conversation_context: Previous conversation context for embedding
        question_embedding: Pre-normalized embedding vector (np.ndarray) for the question
        context_embedding: Pre-normalized embedding vector (np.ndarray) for the context
        sql_query: The SQL query to re-execute on cache hit
        description: Description of the query
        conversation_id: The conversation ID where this query originated
        created_at: When the entry was created
        last_accessed_at: Last access time for LRU eviction
        message_id: The original Genie message ID (for feedback on cache hits)
    """

    genie_space_id: str
    question: str
    conversation_context: str
    question_embedding: np.ndarray
    context_embedding: np.ndarray
    sql_query: str
    description: str
    conversation_id: str
    created_at: datetime
    last_accessed_at: datetime  # Track last access time for LRU eviction
    message_id: str | None = None  # Original Genie message ID for feedback

    def __post_init__(self) -> None:
        """Auto-normalize embeddings to unit vectors for fast dot-product similarity."""
        if not isinstance(self.question_embedding, np.ndarray):
            self.question_embedding = np.array(self.question_embedding)
        self.question_embedding = _normalize(self.question_embedding)

        if not isinstance(self.context_embedding, np.ndarray):
            self.context_embedding = np.array(self.context_embedding)
        self.context_embedding = _normalize(self.context_embedding)


def _normalize(v: np.ndarray) -> np.ndarray:
    """Normalize a vector to unit length. Returns zero vector if norm is 0."""
    norm = np.linalg.norm(v)
    if norm == 0:
        return v
    return v / norm


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """
    Calculate cosine similarity between two embedding vectors.

    This uses the same distance metric as PostgreSQL pg_vector's <=> operator
    to ensure consistent behavior between in-memory and PostgreSQL caches.
    Cosine similarity = 1 - cosine_distance.

    Args:
        a: First embedding vector
        b: Second embedding vector

    Returns:
        Cosine similarity score where 1.0 = identical, 0 = orthogonal, -1 = opposite
    """
    a_arr = np.array(a)
    b_arr = np.array(b)
    norm_a = np.linalg.norm(a_arr)
    norm_b = np.linalg.norm(b_arr)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a_arr, b_arr) / (norm_a * norm_b))


class InMemoryContextAwareGenieService(ContextAwareGenieService):
    """
    In-memory context-aware caching decorator using dual embeddings for similarity lookup.

    This service caches the SQL query generated by Genie along with dual embeddings
    (question + conversation context) for high-precision semantic matching. On
    subsequent queries, it performs similarity search to find cached queries that
    match both the question intent AND conversation context.

    Cache entries are partitioned by genie_space_id to ensure queries from different
    Genie spaces don't return incorrect cache hits.

    On cache hit, it re-executes the cached SQL using the provided warehouse
    to return fresh data while avoiding the Genie NL-to-SQL translation cost.

    Example:
        from dao_ai.config import GenieInMemoryContextAwareCacheParametersModel
        from dao_ai.genie.cache.context_aware import InMemoryContextAwareGenieService

        cache_params = GenieInMemoryContextAwareCacheParametersModel(
            warehouse=warehouse_model,
            embedding_model="databricks-gte-large-en",
            time_to_live_seconds=86400,  # 24 hours
            similarity_threshold=0.85,
            capacity=1000,  # Limit to 1000 entries
        )
        genie = InMemoryContextAwareGenieService(
            impl=GenieService(Genie(space_id="my-space")),
            parameters=cache_params,
            workspace_client=workspace_client,
        ).initialize()

    Thread-safe: Uses a lock to protect cache operations.
    """

    impl: GenieServiceBase
    parameters: GenieInMemoryContextAwareCacheParametersModel
    _workspace_client: WorkspaceClient | None
    name: str
    _embeddings: Any  # DatabricksEmbeddings
    _cache: list[InMemoryCacheEntry]
    _lock: Lock
    _embedding_dims: int | None
    _setup_complete: bool

    def __init__(
        self,
        impl: GenieServiceBase,
        parameters: GenieInMemoryContextAwareCacheParametersModel,
        workspace_client: WorkspaceClient | None = None,
        name: str | None = None,
    ) -> None:
        """
        Initialize the in-memory context-aware cache service.

        Args:
            impl: The underlying GenieServiceBase to delegate to on cache miss.
                The space_id will be obtained from impl.space_id.
            parameters: Cache configuration including warehouse, embedding model, and thresholds
            workspace_client: Optional WorkspaceClient for retrieving conversation history.
                If None, conversation context will not be used.
            name: Name for this cache layer (for logging). Defaults to class name.
        """
        self.impl = impl
        self.parameters = parameters
        self._workspace_client = workspace_client
        self.name = name if name is not None else self.__class__.__name__
        self._embeddings = None
        self._cache = []
        self._lock = Lock()
        self._embedding_dims = None
        self._setup_complete = False

    def _setup(self) -> None:
        """Initialize embeddings model lazily."""
        if self._setup_complete:
            return

        # Initialize embeddings using base class helper
        self._initialize_embeddings(
            self.parameters.embedding_model,
            self.parameters.embedding_dims,
        )

        self._setup_complete = True
        logger.debug(
            "In-memory context-aware cache initialized",
            layer=self.name,
            space_id=self.space_id,
            dims=self._embedding_dims,
            capacity=self.parameters.capacity,
        )

    # Property implementations
    @property
    def warehouse(self) -> WarehouseModel:
        """The warehouse used for executing cached SQL queries."""
        return self.parameters.warehouse

    @property
    def time_to_live(self) -> timedelta | None:
        """Time-to-live for cache entries. None means never expires."""
        ttl = self.parameters.time_to_live_seconds
        if ttl is None or ttl < 0:
            return None
        return timedelta(seconds=ttl)

    @property
    def similarity_threshold(self) -> float:
        """Minimum similarity for cache hit (using L2 distance converted to similarity)."""
        return self.parameters.similarity_threshold

    @property
    def context_similarity_threshold(self) -> float:
        """Minimum similarity for context matching."""
        return self.parameters.context_similarity_threshold

    @property
    def question_weight(self) -> float:
        """Weight for question similarity in combined score."""
        return self.parameters.question_weight

    @property
    def context_weight(self) -> float:
        """Weight for context similarity in combined score."""
        return self.parameters.context_weight

    def _embed_question(
        self, question: str, conversation_id: str | None = None
    ) -> tuple[list[float], list[float], str]:
        """
        Generate dual embeddings using Genie API for conversation history.

        Args:
            question: The question to embed
            conversation_id: Optional conversation ID for retrieving context

        Returns:
            Tuple of (question_embedding, context_embedding, conversation_context_string)
        """
        return self._embed_question_with_genie_history(
            question,
            conversation_id,
            self.parameters.context_window_size,
            self.parameters.max_context_tokens,
        )

    @mlflow.trace(name="semantic_search_in_memory")
    def _find_similar(
        self,
        question: str,
        conversation_context: str,
        question_embedding: list[float],
        context_embedding: list[float],
        conversation_id: str | None = None,
    ) -> tuple[SQLCacheEntry, float] | None:
        """
        Find a semantically similar cached entry using dual embedding matching.

        Performs linear scan through cached entries for this space, computing
        cosine similarity via dot product on pre-normalized embeddings.

        Lock is held only briefly to snapshot entries and clean up expired ones.
        Similarity computation happens outside the lock to allow concurrent reads.

        Args:
            question: The original question (for logging)
            conversation_context: The conversation context string
            question_embedding: The embedding vector of just the question
            context_embedding: The embedding vector of the conversation context
            conversation_id: Optional conversation ID (for logging)

        Returns:
            Tuple of (SQLCacheEntry, combined_similarity_score) if found, None otherwise
        """
        ttl_seconds = self.parameters.time_to_live_seconds
        ttl_disabled = ttl_seconds is None or ttl_seconds < 0

        question_weight = self.question_weight
        context_weight = self.context_weight

        # Pre-normalize query embeddings once for fast dot-product similarity
        q_norm = _normalize(np.array(question_embedding))
        c_norm = _normalize(np.array(context_embedding))

        # Cache datetime.now() once for all TTL checks
        now = datetime.now()

        # Step 1: Snapshot space entries and identify expired ones under lock
        expired_ids: set[int] = set()
        with self._lock:
            space_entries = [
                entry for entry in self._cache if entry.genie_space_id == self.space_id
            ]

        # Step 2: Compute similarities outside the lock (read-only on entries)
        valid_entries: list[InMemoryCacheEntry] = []
        for entry in space_entries:
            if not ttl_disabled:
                age = now - entry.created_at
                if age.total_seconds() > ttl_seconds:
                    expired_ids.add(id(entry))
                    continue
            valid_entries.append(entry)

        best_entry: InMemoryCacheEntry | None = None
        best_question_sim: float = 0.0
        best_context_sim: float = 0.0
        best_combined_sim: float = 0.0

        similarity_threshold = self.similarity_threshold
        for entry in valid_entries:
            # Dot product of pre-normalized vectors = cosine similarity
            question_sim = float(np.dot(q_norm, entry.question_embedding))

            # Early-exit: skip entries that can't meet the question threshold
            if question_sim < similarity_threshold:
                continue

            context_sim = float(np.dot(c_norm, entry.context_embedding))

            combined_sim = (question_weight * question_sim) + (
                context_weight * context_sim
            )

            if combined_sim > best_combined_sim:
                best_entry = entry
                best_question_sim = question_sim
                best_context_sim = context_sim
                best_combined_sim = combined_sim

        # Step 3: Clean up expired entries under lock
        if expired_ids:
            with self._lock:
                self._cache = [e for e in self._cache if id(e) not in expired_ids]
                logger.trace(
                    "Deleted expired entries",
                    layer=self.name,
                    count=len(expired_ids),
                )

        # No entries found
        if best_entry is None:
            logger.info(
                "Cache MISS (no entries)",
                layer=self.name,
                question=question[:50],
                space=self.space_id,
                delegating_to=type(self.impl).__name__,
            )
            return None

        # Log best match info
        logger.debug(
            "Best match found",
            layer=self.name,
            question_sim=f"{best_question_sim:.4f}",
            context_sim=f"{best_context_sim:.4f}",
            combined_sim=f"{best_combined_sim:.4f}",
            cached_question=best_entry.question[:50],
        )

        # Check question similarity threshold
        if best_question_sim < self.similarity_threshold:
            logger.info(
                "Cache MISS (question similarity too low)",
                layer=self.name,
                question_sim=f"{best_question_sim:.4f}",
                threshold=self.similarity_threshold,
            )
            return None

        # Bidirectional context-skip: skip context threshold when
        # EITHER side has no context. This handles two scenarios:
        # 1. Cached entry had no context (stored without conversation history)
        # 2. Query has no context (first message in a new conversation)
        # In both cases the entry is "context-agnostic" and should
        # match on question similarity alone.
        has_stored_context = bool(
            best_entry.conversation_context and best_entry.conversation_context.strip()
        )
        query_has_context = bool(conversation_context and conversation_context.strip())

        if has_stored_context and query_has_context:
            if best_context_sim < self.context_similarity_threshold:
                logger.info(
                    "Cache MISS (context similarity too low)",
                    layer=self.name,
                    context_sim=f"{best_context_sim:.4f}",
                    threshold=self.context_similarity_threshold,
                )
                return None
        else:
            logger.debug(
                "Context threshold skipped (no context on one or both sides)",
                layer=self.name,
                context_sim=f"{best_context_sim:.4f}",
                threshold=self.context_similarity_threshold,
                has_stored_context=has_stored_context,
                query_has_context=query_has_context,
            )

        # Cache HIT - Update last accessed time
        with self._lock:
            best_entry.last_accessed_at = now

        cache_age_seconds = (now - best_entry.created_at).total_seconds()
        logger.info(
            "Cache HIT",
            layer=self.name,
            question=question[:80],
            conversation_id=conversation_id,
            matched_question=best_entry.question[:80],
            cache_age_seconds=round(cache_age_seconds, 1),
            question_similarity=f"{best_question_sim:.4f}",
            context_similarity=f"{best_context_sim:.4f}",
            combined_similarity=f"{best_combined_sim:.4f}",
        )

        cache_entry = SQLCacheEntry(
            query=best_entry.sql_query,
            description=best_entry.description,
            conversation_id=best_entry.conversation_id,
            created_at=best_entry.created_at,
            message_id=best_entry.message_id,
            # In-memory caches don't have database row IDs
            cache_entry_id=None,
        )
        return cache_entry, best_combined_sim

    def _store_entry(
        self,
        question: str,
        conversation_context: str,
        question_embedding: list[float],
        context_embedding: list[float],
        response: GenieResponse,
        message_id: str | None = None,
    ) -> None:
        """
        Store a new cache entry with pre-normalized dual embeddings and message_id.

        Embeddings are converted to NumPy arrays and normalized at storage time
        so that cosine similarity can be computed as a simple dot product at
        lookup time.

        If capacity is set and reached, evicts least recently used entries in a
        single O(n) pass instead of repeated scans.
        """
        now = datetime.now()
        new_entry = InMemoryCacheEntry(
            genie_space_id=self.space_id,
            question=question,
            conversation_context=conversation_context,
            question_embedding=question_embedding,
            context_embedding=context_embedding,
            sql_query=response.query,
            description=response.description,
            conversation_id=response.conversation_id,
            created_at=now,
            last_accessed_at=now,
            message_id=message_id,
        )

        with self._lock:
            # Enforce capacity limit (LRU eviction) in a single pass
            if self.parameters.capacity is not None:
                space_count = sum(
                    1 for e in self._cache if e.genie_space_id == self.space_id
                )
                # +1 because we're about to add a new entry
                entries_to_evict = max(0, space_count - self.parameters.capacity + 1)

                if entries_to_evict > 0:
                    # Collect space entries sorted by LRU (oldest access first)
                    space_entries = sorted(
                        (e for e in self._cache if e.genie_space_id == self.space_id),
                        key=lambda e: e.last_accessed_at,
                    )
                    evict_ids = {
                        id(space_entries[i])
                        for i in range(min(entries_to_evict, len(space_entries)))
                    }
                    for evicted in space_entries[:entries_to_evict]:
                        logger.trace(
                            "Evicted LRU cache entry",
                            layer=self.name,
                            question=evicted.question[:50],
                            capacity=self.parameters.capacity,
                        )
                    self._cache = [e for e in self._cache if id(e) not in evict_ids]

            self._cache.append(new_entry)
            new_size = sum(1 for e in self._cache if e.genie_space_id == self.space_id)
            logger.debug(
                "Stored cache entry",
                layer=self.name,
                question=question[:50],
                space=self.space_id,
                cache_size=new_size,
                capacity=self.parameters.capacity,
            )

    def _on_stale_cache_entry(self, question: str) -> None:
        """Remove stale cache entry from memory."""
        with self._lock:
            for idx, entry in enumerate(self._cache):
                if entry.genie_space_id == self.space_id and entry.question == question:
                    del self._cache[idx]
                    logger.info(
                        "Deleted stale cache entry from memory",
                        layer=self.name,
                        question=question[:50],
                    )
                    break

    def _invalidate_by_question(self, question: str) -> bool:
        """
        Invalidate cache entries matching a specific question.

        This method is called when negative feedback is received to remove
        the corresponding cache entry from the in-memory cache.

        Args:
            question: The question text to match and invalidate

        Returns:
            True if an entry was found and invalidated, False otherwise
        """
        with self._lock:
            for idx, entry in enumerate(self._cache):
                if entry.genie_space_id == self.space_id and entry.question == question:
                    del self._cache[idx]
                    logger.info(
                        "Invalidated cache entry by question",
                        layer=self.name,
                        question=question[:50],
                        space_id=self.space_id,
                    )
                    return True
            return False

    # Note: ask_question_with_cache_info is inherited from ContextAwareGenieService
    # using the Template Method pattern. InMemoryContextAwareGenieService uses the
    # default empty hook implementations since it doesn't track prompt history.

    # Template Method implementations for invalidate_expired() and clear()

    def _delete_expired_entries(self, ttl_seconds: int) -> int:
        """Delete expired entries from the cache."""
        deleted = 0
        with self._lock:
            indices_to_delete: list[int] = []
            now = datetime.now()

            for idx, entry in enumerate(self._cache):
                if entry.genie_space_id != self.space_id:
                    continue

                age = now - entry.created_at
                if age.total_seconds() > ttl_seconds:
                    indices_to_delete.append(idx)

            # Delete in reverse order
            for idx in reversed(indices_to_delete):
                del self._cache[idx]
                deleted += 1

            logger.trace(
                "Deleted expired entries",
                layer=self.name,
                deleted_count=deleted,
            )

        return deleted

    def _delete_all_entries(self) -> int:
        """Delete all entries for this Genie space."""
        deleted = 0

        with self._lock:
            indices_to_delete: list[int] = []
            for idx, entry in enumerate(self._cache):
                if entry.genie_space_id == self.space_id:
                    indices_to_delete.append(idx)

            # Delete in reverse order
            for idx in reversed(indices_to_delete):
                del self._cache[idx]
                deleted += 1

            logger.debug(
                "Cleared cache entries", layer=self.name, deleted_count=deleted
            )

        return deleted

    @property
    def size(self) -> int:
        """Current number of entries in the cache for this Genie space."""
        self._setup()
        with self._lock:
            return sum(1 for e in self._cache if e.genie_space_id == self.space_id)

    # Template Method implementations for stats()

    def _count_all_entries(self) -> int:
        """Count all cache entries for this Genie space."""
        with self._lock:
            return sum(1 for e in self._cache if e.genie_space_id == self.space_id)

    def _count_entries_with_ttl(self, ttl_seconds: int) -> tuple[int, int]:
        """Count total and expired entries for this Genie space."""
        now = datetime.now()
        with self._lock:
            space_entries = [
                e for e in self._cache if e.genie_space_id == self.space_id
            ]
            total = len(space_entries)
            expired = 0
            for entry in space_entries:
                age = now - entry.created_at
                if age.total_seconds() > ttl_seconds:
                    expired += 1
            return total, expired

    def _get_additional_stats(self) -> dict[str, Any]:
        """Add capacity info to stats."""
        return {"capacity": self.parameters.capacity}

    def get_entries(
        self,
        limit: int | None = None,
        offset: int | None = None,
        include_embeddings: bool = False,
        conversation_id: str | None = None,
        created_after: datetime | None = None,
        created_before: datetime | None = None,
        question_contains: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Get cache entries with optional filtering.

        This method retrieves cache entries for inspection, debugging, or
        generating evaluation datasets for threshold optimization.

        Args:
            limit: Maximum number of entries to return (None = no limit)
            offset: Number of entries to skip for pagination (None = 0)
            include_embeddings: Whether to include embedding vectors in results.
                Embeddings are large, so set False for general inspection.
            conversation_id: Filter by conversation ID (None = all conversations)
            created_after: Only entries created after this time (None = no filter)
            created_before: Only entries created before this time (None = no filter)
            question_contains: Case-insensitive text search on question field

        Returns:
            List of cache entry dicts. See base class for full key documentation.

        Example:
            # Get entries with embeddings for evaluation dataset generation
            entries = cache.get_entries(include_embeddings=True, limit=100)
            eval_dataset = generate_eval_dataset_from_cache(entries)
        """
        self._setup()

        with self._lock:
            # Filter entries for this space
            filtered_entries: list[InMemoryCacheEntry] = []

            for entry in self._cache:
                # Filter by space_id
                if entry.genie_space_id != self.space_id:
                    continue

                # Filter by conversation_id
                if (
                    conversation_id is not None
                    and entry.conversation_id != conversation_id
                ):
                    continue

                # Filter by created_after
                if created_after is not None and entry.created_at <= created_after:
                    continue

                # Filter by created_before
                if created_before is not None and entry.created_at >= created_before:
                    continue

                # Filter by question_contains (case-insensitive)
                if question_contains is not None:
                    if question_contains.lower() not in entry.question.lower():
                        continue

                filtered_entries.append(entry)

            # Sort by created_at descending (most recent first)
            filtered_entries.sort(key=lambda e: e.created_at, reverse=True)

            # Apply offset
            if offset is not None and offset > 0:
                filtered_entries = filtered_entries[offset:]

            # Apply limit
            if limit is not None:
                filtered_entries = filtered_entries[:limit]

            # Convert to dicts
            entries: list[dict[str, Any]] = []
            for entry in filtered_entries:
                result: dict[str, Any] = {
                    "id": None,  # In-memory caches don't have database IDs
                    "question": entry.question,
                    "conversation_context": entry.conversation_context,
                    "sql_query": entry.sql_query,
                    "description": entry.description,
                    "conversation_id": entry.conversation_id,
                    "created_at": entry.created_at,
                }

                if include_embeddings:
                    result["question_embedding"] = entry.question_embedding.tolist()
                    result["context_embedding"] = entry.context_embedding.tolist()

                entries.append(result)

            logger.debug(
                "Retrieved cache entries",
                layer=self.name,
                count=len(entries),
                include_embeddings=include_embeddings,
                filters={
                    "conversation_id": conversation_id,
                    "created_after": str(created_after) if created_after else None,
                    "created_before": str(created_before) if created_before else None,
                    "question_contains": question_contains,
                },
            )

            return entries
