"""
LRU (Least Recently Used) cache implementation for Genie SQL queries.

This module provides an in-memory LRU cache that stores SQL queries generated
by Genie. On cache hit, the cached SQL is re-executed against the warehouse
to return fresh data while avoiding the Genie NL-to-SQL translation cost.
"""

from __future__ import annotations

from collections import OrderedDict
from datetime import datetime, timedelta
from threading import Lock

import mlflow
import pandas as pd
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.dashboards import GenieFeedbackRating
from loguru import logger

from dao_ai.config import GenieLRUCacheParametersModel, WarehouseModel
from dao_ai.genie.cache.base import (
    CacheResult,
    GenieServiceBase,
    SQLCacheEntry,
)
from dao_ai.genie.cache.core import execute_sql_via_warehouse
from dao_ai.genie.core import GenieResponse


class LRUCacheService(GenieServiceBase):
    """
    LRU caching decorator that caches SQL queries and re-executes them.

    This service caches the SQL query generated by Genie (not the result data).
    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 GenieLRUCacheParametersModel, WarehouseModel
        from dao_ai.genie.cache import LRUCacheService

        cache_params = GenieLRUCacheParametersModel(
            warehouse=warehouse_model,
            capacity=100,
            time_to_live_seconds=86400  # 24 hours
        )
        genie = LRUCacheService(
            impl=GenieService(Genie(space_id="my-space")),
            parameters=cache_params
        )

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

    impl: GenieServiceBase
    parameters: GenieLRUCacheParametersModel
    name: str
    _cache: OrderedDict[str, SQLCacheEntry]
    _lock: Lock

    def __init__(
        self,
        impl: GenieServiceBase,
        parameters: GenieLRUCacheParametersModel,
        name: str | None = None,
    ) -> None:
        """
        Initialize the SQL cache service.

        Args:
            impl: The underlying GenieServiceBase to delegate to on cache miss
            parameters: Cache configuration including warehouse, capacity, and TTL
            name: Name for this cache layer (for logging). Defaults to class name.
        """
        self.impl = impl
        self.parameters = parameters
        self.name = name if name is not None else self.__class__.__name__
        self._cache = OrderedDict()
        self._lock = Lock()

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

    @property
    def capacity(self) -> int:
        """Maximum number of SQL queries to cache."""
        return self.parameters.capacity

    @property
    def time_to_live(self) -> timedelta | None:
        """Duration after which cached queries expire. None means never expires."""
        ttl = self.parameters.time_to_live_seconds
        if ttl is None or ttl < 0:
            return None
        return timedelta(seconds=ttl)

    @staticmethod
    def _normalize_key(question: str, conversation_id: str | None = None) -> str:
        """
        Normalize the question and conversation_id to create a consistent cache key.

        Args:
            question: The question text
            conversation_id: Optional conversation ID to include in the key

        Returns:
            A normalized cache key combining question and conversation_id
        """
        normalized_question = question.strip().lower()
        if conversation_id:
            return f"{conversation_id}::{normalized_question}"
        return normalized_question

    def _is_expired(self, entry: SQLCacheEntry) -> bool:
        """Check if a cache entry has exceeded its TTL. Returns False if TTL is disabled."""
        if self.time_to_live is None:
            return False
        age: timedelta = datetime.now() - entry.created_at
        return age > self.time_to_live

    def _evict_oldest(self) -> None:
        """Remove the oldest (least recently used) entry."""
        if self._cache:
            oldest_key: str = next(iter(self._cache))
            del self._cache[oldest_key]
            logger.trace("Evicted cache entry", layer=self.name, key=oldest_key[:50])

    def _get(self, key: str) -> SQLCacheEntry | None:
        """Get from cache, returning None if not found or expired."""
        if key not in self._cache:
            return None

        entry: SQLCacheEntry = self._cache[key]

        if self._is_expired(entry):
            del self._cache[key]
            logger.trace("Expired cache entry", layer=self.name, key=key[:50])
            return None

        self._cache.move_to_end(key)
        return entry

    def _put(
        self, key: str, response: GenieResponse, message_id: str | None = None
    ) -> None:
        """Store SQL query in cache, evicting if at capacity."""
        # Skip caching if query is empty or whitespace
        if not response.query or not response.query.strip():
            logger.warning(
                "Not caching: response has no SQL query",
                layer=self.name,
                key=key[:50],
                description=response.description[:80] if response.description else None,
            )
            return

        if key in self._cache:
            del self._cache[key]

        while len(self._cache) >= self.capacity:
            self._evict_oldest()

        self._cache[key] = SQLCacheEntry(
            query=response.query,
            description=response.description,
            conversation_id=response.conversation_id,
            created_at=datetime.now(),
            message_id=message_id,
            # LRU cache is in-memory only, no database row ID
            cache_entry_id=None,
        )
        logger.debug(
            "Stored cache entry",
            layer=self.name,
            key=key[:50],
            sql=response.query[:50] if response.query else None,
            cache_size=len(self._cache),
            capacity=self.capacity,
            message_id=message_id,
        )

    @mlflow.trace(name="execute_cached_sql")
    def _execute_sql(self, sql: str) -> pd.DataFrame | str:
        """
        Execute SQL using the warehouse and return results as DataFrame.

        Args:
            sql: The SQL query to execute

        Returns:
            DataFrame with results, or error message string
        """
        # Validate SQL is not empty
        if not sql or not sql.strip():
            error_msg: str = "Cannot execute empty SQL query"
            logger.error(
                "SQL execution failed: empty query",
                layer=self.name,
                sql=repr(sql),
            )
            return error_msg

        # Use shared utility function for SQL execution
        return execute_sql_via_warehouse(
            warehouse=self.warehouse,
            sql=sql,
            layer_name=self.name,
        )

    def ask_question(
        self, question: str, conversation_id: str | None = None
    ) -> CacheResult:
        """
        Ask a question, using cached SQL query if available.

        On cache hit, re-executes the cached SQL to get fresh data.
        Returns CacheResult with cache metadata.
        """
        return self.ask_question_with_cache_info(question, conversation_id)

    @mlflow.trace(name="genie_lru_cache_lookup")
    def ask_question_with_cache_info(
        self,
        question: str,
        conversation_id: str | None = None,
    ) -> CacheResult:
        """
        Ask a question with detailed cache hit information.

        On cache hit, the cached SQL is re-executed to return fresh data.

        Args:
            question: The question to ask
            conversation_id: Optional conversation ID

        Returns:
            CacheResult with fresh response and cache metadata
        """
        key: str = self._normalize_key(question, conversation_id)

        # Check cache
        with self._lock:
            cached: SQLCacheEntry | None = self._get(key)

        if cached is not None:
            # Defensive check: if cached query is empty, treat as cache miss
            if not cached.query or not cached.query.strip():
                logger.warning(
                    "Cache HIT but query is empty, treating as MISS",
                    layer=self.name,
                    question=question[:80],
                    conversation_id=conversation_id,
                    key=key[:50],
                )
                # Invalidate this bad cache entry
                with self._lock:
                    if key in self._cache:
                        del self._cache[key]
                # Fall through to cache miss logic below
            else:
                cache_age_seconds = (datetime.now() - cached.created_at).total_seconds()
                logger.info(
                    "Cache HIT",
                    layer=self.name,
                    question=question[:80],
                    conversation_id=conversation_id,
                    cached_sql=cached.query[:80] if cached.query else None,
                    cache_age_seconds=round(cache_age_seconds, 1),
                    cache_size=self.size,
                    capacity=self.capacity,
                    ttl_seconds=self.parameters.time_to_live_seconds,
                )

                # Re-execute the cached SQL to get fresh data
                result: pd.DataFrame | str = self._execute_sql(cached.query)

                # Check if SQL execution failed (returns error string instead of DataFrame)
                if isinstance(result, str):
                    logger.warning(
                        "Cached SQL execution failed, falling back to Genie",
                        layer=self.name,
                        question=question[:80],
                        conversation_id=conversation_id,
                        cached_sql=cached.query[:80],
                        error=result[:200],
                        cache_key=key[:50],
                    )

                    # Invalidate the bad cache entry
                    with self._lock:
                        if key in self._cache:
                            del self._cache[key]
                            logger.info(
                                "Invalidated stale cache entry",
                                layer=self.name,
                                cache_key=key[:50],
                                cache_size=len(self._cache),
                                capacity=self.capacity,
                            )

                    # Fall back to Genie to get fresh SQL
                    logger.info(
                        "Delegating to Genie for fresh SQL",
                        layer=self.name,
                        question=question[:80],
                        delegating_to=type(self.impl).__name__,
                    )
                    fallback_result: CacheResult = self.impl.ask_question(
                        question, conversation_id
                    )

                    # Store the fresh SQL in cache (including message_id for feedback)
                    if fallback_result.response.query:
                        with self._lock:
                            self._put(
                                key,
                                fallback_result.response,
                                message_id=fallback_result.message_id,
                            )
                        logger.info(
                            "Stored fresh SQL from fallback",
                            layer=self.name,
                            fresh_sql=fallback_result.response.query[:80],
                            cache_size=len(self._cache),
                            capacity=self.capacity,
                            message_id=fallback_result.message_id,
                        )
                    else:
                        logger.warning(
                            "Fallback response has no SQL query to cache",
                            layer=self.name,
                            question=question[:80],
                        )

                    logger.info(
                        "Fallback completed successfully",
                        layer=self.name,
                        question=question[:80],
                        fallback_from="stale_cache",
                        has_result=fallback_result.response.result is not None,
                    )

                    # Return as cache miss (fallback scenario)
                    # Propagate message_id from fallback result
                    return CacheResult(
                        response=fallback_result.response,
                        cache_hit=False,
                        served_by=None,
                        message_id=fallback_result.message_id,
                    )

                # Use current conversation_id, not the cached one
                response: GenieResponse = GenieResponse(
                    result=result,
                    query=cached.query,
                    description=cached.description,
                    conversation_id=conversation_id
                    if conversation_id
                    else cached.conversation_id,
                )

                # Cache hit - include message_id from original response for feedback support
                return CacheResult(
                    response=response,
                    cache_hit=True,
                    served_by=self.name,
                    message_id=cached.message_id,
                    # LRU cache is in-memory only, no cache_entry_id for traceability
                    cache_entry_id=None,
                )

        # Cache miss - delegate to wrapped service
        logger.info(
            "Cache MISS",
            layer=self.name,
            question=question[:80],
            conversation_id=conversation_id,
            cache_size=self.size,
            capacity=self.capacity,
            ttl_seconds=self.parameters.time_to_live_seconds,
            delegating_to=type(self.impl).__name__,
        )

        result: CacheResult = self.impl.ask_question(question, conversation_id)
        with self._lock:
            self._put(key, result.response, message_id=result.message_id)
        # Propagate the inner cache's result - if it was a hit there, preserve that info
        return result

    @property
    def space_id(self) -> str:
        return self.impl.space_id

    @property
    def workspace_client(self) -> WorkspaceClient | None:
        """Get workspace client by delegating to impl."""
        return self.impl.workspace_client

    def invalidate(self, question: str, conversation_id: str | None = None) -> bool:
        """
        Remove a specific entry from the cache.

        Args:
            question: The question text
            conversation_id: Optional conversation ID to match

        Returns:
            True if the entry was found and removed, False otherwise
        """
        key: str = self._normalize_key(question, conversation_id)
        with self._lock:
            if key in self._cache:
                del self._cache[key]
                return True
            return False

    def clear(self) -> int:
        """Clear all entries from the cache."""
        with self._lock:
            count: int = len(self._cache)
            self._cache.clear()
            return count

    @property
    def size(self) -> int:
        """Current number of entries in the cache."""
        with self._lock:
            return len(self._cache)

    def stats(self) -> dict[str, int | float | None]:
        """Return cache statistics."""
        with self._lock:
            expired: int = sum(1 for e in self._cache.values() if self._is_expired(e))
            ttl = self.time_to_live
            return {
                "size": len(self._cache),
                "capacity": self.capacity,
                "ttl_seconds": ttl.total_seconds() if ttl else None,
                "expired_entries": expired,
                "valid_entries": len(self._cache) - expired,
            }

    @mlflow.trace(name="genie_lru_cache_send_feedback")
    def send_feedback(
        self,
        conversation_id: str,
        rating: GenieFeedbackRating,
        message_id: str | None = None,
        was_cache_hit: bool = False,
    ) -> None:
        """
        Send feedback for a Genie message with cache invalidation.

        For LRU cache, this method:
        1. If was_cache_hit is False: forwards feedback to the underlying service
        2. If rating is NEGATIVE: invalidates any matching cache entries

        Args:
            conversation_id: The conversation containing the message
            rating: The feedback rating (POSITIVE, NEGATIVE, or NONE)
            message_id: Optional message ID. If None, looks up the most recent message.
            was_cache_hit: Whether the response being rated was served from cache.

        Note:
            For cached responses (was_cache_hit=True), only cache invalidation is
            performed. No feedback is sent to the Genie API because cached responses
            don't have a corresponding Genie message.

            Future Enhancement: To enable full Genie feedback for cached responses,
            the cache would need to store the original message_id. See GenieServiceBase
            docstring for details on required changes.
        """
        # Handle cache invalidation on negative feedback
        invalidated = False
        if rating == GenieFeedbackRating.NEGATIVE:
            # For LRU cache, we invalidate by conversation_id since that's part of the key
            # Iterate through cache and remove entries matching the conversation_id
            with self._lock:
                keys_to_remove: list[str] = []
                for key, entry in self._cache.items():
                    if entry.conversation_id == conversation_id:
                        keys_to_remove.append(key)

                for key in keys_to_remove:
                    del self._cache[key]
                    invalidated = True
                    logger.info(
                        "Invalidated cache entry due to negative feedback",
                        layer=self.name,
                        cache_key=key[:50],
                        conversation_id=conversation_id,
                    )

            if not keys_to_remove:
                logger.debug(
                    "No cache entries found to invalidate for negative feedback",
                    layer=self.name,
                    conversation_id=conversation_id,
                )

        # Forward feedback to underlying service if not a cache hit
        # For cache hits, there's no Genie message to provide feedback on
        if was_cache_hit:
            logger.info(
                "Skipping Genie API feedback - response was served from cache",
                layer=self.name,
                conversation_id=conversation_id,
                rating=rating.value if rating else None,
                cache_invalidated=invalidated,
            )
            return

        # Forward to underlying service
        logger.debug(
            "Forwarding feedback to underlying service",
            layer=self.name,
            conversation_id=conversation_id,
            rating=rating.value if rating else None,
            delegating_to=type(self.impl).__name__,
        )
        self.impl.send_feedback(
            conversation_id=conversation_id,
            rating=rating,
            message_id=message_id,
            was_cache_hit=False,  # Already handled, so pass False
        )
