"""
Semantic caching for Natural Language Query system.

Uses Redis with embedding similarity matching for intelligent caching.
Caches both query results and LLM responses.
"""

from typing import Any, Dict, Optional, Tuple
import json
import hashlib
import os


class CacheManager:
    """
    Semantic cache using Redis and embedding similarity.

    Caches query results based on semantic similarity rather than exact matches.
    Uses sentence embeddings to find similar queries.
    """

    def __init__(
        self,
        redis_url: Optional[str] = None,
        ttl: int = 3600,
        similarity_threshold: float = 0.85,
        use_redis: bool = True,
    ):
        """
        Initialize cache manager.

        Args:
            redis_url: Redis connection URL (default: from env or localhost)
            ttl: Time-to-live for cache entries in seconds
            similarity_threshold: Minimum similarity score for cache hits (0-1)
            use_redis: If False, use in-memory cache (for testing)
        """
        self.ttl = ttl
        self.similarity_threshold = similarity_threshold
        self.use_redis = use_redis

        # In-memory fallback cache
        self._memory_cache: Dict[str, Any] = {}

        if use_redis:
            try:
                import redis
                from redis import Redis

                if redis_url is None:
                    redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")

                self.redis_client: Optional[Redis] = redis.from_url(
                    redis_url, decode_responses=True
                )

                # Test connection
                self.redis_client.ping()
                self.redis_available = True
                print("Redis cache connected successfully")

            except (ImportError, Exception) as e:
                print(
                    f"Redis not available, using in-memory cache: {e}"
                )
                self.redis_client = None
                self.redis_available = False
        else:
            self.redis_client = None
            self.redis_available = False

    def _compute_query_hash(self, query: str) -> str:
        """
        Compute hash for query (for exact matching).

        Args:
            query: User query

        Returns:
            str: Hash string
        """
        return hashlib.sha256(query.lower().strip().encode()).hexdigest()

    def _get_embedding(self, text: str) -> list:
        """
        Get text embedding for semantic similarity.

        Args:
            text: Text to embed

        Returns:
            list: Embedding vector
        """
        try:
            from sentence_transformers import SentenceTransformer

            # Use a lightweight model
            model = SentenceTransformer("all-MiniLM-L6-v2")
            embedding = model.encode(text, convert_to_tensor=False)
            return embedding.tolist()

        except ImportError:
            # Fallback to hash-based matching if sentence-transformers not available
            return []

    def _cosine_similarity(self, vec1: list, vec2: list) -> float:
        """
        Calculate cosine similarity between two vectors.

        Args:
            vec1: First vector
            vec2: Second vector

        Returns:
            float: Similarity score (0-1)
        """
        if not vec1 or not vec2:
            return 0.0

        import numpy as np

        vec1 = np.array(vec1)
        vec2 = np.array(vec2)

        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)

        if norm1 == 0 or norm2 == 0:
            return 0.0

        return float(dot_product / (norm1 * norm2))

    def get(self, query: str) -> Optional[Dict[str, Any]]:
        """
        Get cached result for query.

        Uses semantic similarity to find similar queries.

        Args:
            query: User query

        Returns:
            Optional[dict]: Cached result or None
        """
        query_hash = self._compute_query_hash(query)

        # Try exact match first
        if self.redis_available and self.redis_client:
            try:
                cached = self.redis_client.get(f"nlq:exact:{query_hash}")
                if cached:
                    return json.loads(cached)
            except Exception as e:
                print(f"Redis get error: {e}")

        # Try in-memory cache
        if query_hash in self._memory_cache:
            return self._memory_cache[query_hash]

        # Try semantic similarity match
        if self.redis_available and self.redis_client:
            try:
                query_embedding = self._get_embedding(query)
                if query_embedding:
                    # Search for similar queries
                    cached_queries = self.redis_client.keys("nlq:semantic:*")

                    best_match = None
                    best_similarity = 0.0

                    for key in cached_queries:
                        try:
                            cached_data = json.loads(
                                self.redis_client.get(key)
                            )
                            cached_embedding = cached_data.get("embedding")

                            if cached_embedding:
                                similarity = self._cosine_similarity(
                                    query_embedding, cached_embedding
                                )

                                if (
                                    similarity > best_similarity
                                    and similarity >= self.similarity_threshold
                                ):
                                    best_similarity = similarity
                                    best_match = cached_data.get("result")

                        except Exception:
                            continue

                    if best_match:
                        print(
                            f"Cache hit (semantic similarity: {best_similarity:.2f})"
                        )
                        return best_match

            except Exception as e:
                print(f"Semantic search error: {e}")

        return None

    def set(self, query: str, result: Dict[str, Any]) -> None:
        """
        Cache query result.

        Args:
            query: User query
            result: Query result to cache
        """
        query_hash = self._compute_query_hash(query)

        # Store exact match
        if self.redis_available and self.redis_client:
            try:
                self.redis_client.setex(
                    f"nlq:exact:{query_hash}",
                    self.ttl,
                    json.dumps(result, default=str),
                )

                # Store with embedding for semantic search
                query_embedding = self._get_embedding(query)
                if query_embedding:
                    semantic_data = {
                        "query": query,
                        "embedding": query_embedding,
                        "result": result,
                    }
                    self.redis_client.setex(
                        f"nlq:semantic:{query_hash}",
                        self.ttl,
                        json.dumps(semantic_data, default=str),
                    )

            except Exception as e:
                print(f"Redis set error: {e}")

        # Also store in memory cache as fallback
        self._memory_cache[query_hash] = result

    def clear(self) -> None:
        """Clear all cached data."""
        if self.redis_available and self.redis_client:
            try:
                # Clear all NLQ cache keys
                keys = self.redis_client.keys("nlq:*")
                if keys:
                    self.redis_client.delete(*keys)
            except Exception as e:
                print(f"Redis clear error: {e}")

        # Clear memory cache
        self._memory_cache.clear()

    def get_stats(self) -> Dict[str, Any]:
        """
        Get cache statistics.

        Returns:
            dict: Cache stats
        """
        stats = {
            "redis_available": self.redis_available,
            "memory_cache_size": len(self._memory_cache),
            "ttl": self.ttl,
            "similarity_threshold": self.similarity_threshold,
        }

        if self.redis_available and self.redis_client:
            try:
                exact_keys = len(self.redis_client.keys("nlq:exact:*"))
                semantic_keys = len(self.redis_client.keys("nlq:semantic:*"))

                stats["redis_exact_cache_size"] = exact_keys
                stats["redis_semantic_cache_size"] = semantic_keys
                stats["redis_total_size"] = exact_keys + semantic_keys

            except Exception as e:
                stats["redis_error"] = str(e)

        return stats


def test_cache():
    """Test the cache manager."""
    cache = CacheManager(use_redis=False)

    # Test 1: Store and retrieve
    print("Test 1: Store and retrieve")
    query1 = "What is the conversion rate?"
    result1 = {"answer": "3.5%", "code": "..."}

    cache.set(query1, result1)
    retrieved = cache.get(query1)

    print(f"Stored: {result1}")
    print(f"Retrieved: {retrieved}")
    print(f"Match: {retrieved == result1}")
    print()

    # Test 2: Similar query (semantic matching)
    print("Test 2: Semantic matching")
    query2 = "What's the conversion percentage?"  # Similar to query1
    retrieved2 = cache.get(query2)

    print(f"Original query: {query1}")
    print(f"Similar query: {query2}")
    print(f"Retrieved: {retrieved2}")
    print()

    # Test 3: Stats
    print("Test 3: Cache stats")
    stats = cache.get_stats()
    print(json.dumps(stats, indent=2))


if __name__ == "__main__":
    test_cache()
