"""
semantic_cache.py
-----------------
Semantic caching for LLM responses.

Uses HIEROMEM's vector index to cache and retrieve LLM responses
based on semantic similarity, reducing API costs and latency.

Author: Antigravity
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Callable, Dict, List, Optional, Tuple
import numpy as np


@dataclass
class CacheEntry:
    """A cached LLM response."""
    prompt: str
    response: str
    embedding: np.ndarray
    created_at: datetime
    ttl_seconds: Optional[int]
    hits: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def is_expired(self) -> bool:
        if self.ttl_seconds is None:
            return False
        return datetime.now() > self.created_at + timedelta(seconds=self.ttl_seconds)


@dataclass
class CacheStats:
    """Cache statistics."""
    total_requests: int = 0
    cache_hits: int = 0
    cache_misses: int = 0
    total_saved_ms: float = 0.0

    @property
    def hit_rate(self) -> float:
        if self.total_requests == 0:
            return 0.0
        return self.cache_hits / self.total_requests

    def to_dict(self) -> Dict[str, Any]:
        return {
            "total_requests": self.total_requests,
            "cache_hits": self.cache_hits,
            "cache_misses": self.cache_misses,
            "hit_rate": self.hit_rate,
            "total_saved_ms": self.total_saved_ms
        }


class SemanticCache:
    """
    Semantic cache for LLM responses.

    Uses embedding similarity to find cached responses for similar prompts,
    reducing LLM API calls and latency.

    Example:
        cache = SemanticCache(similarity_threshold=0.95)

        # First call - cache miss, calls LLM
        result = cache.get_or_compute("What is Python?", llm.complete)

        # Second call with similar prompt - cache hit
        result = cache.get_or_compute("Tell me about Python", llm.complete)

        print(cache.stats)  # Shows hit rate, saved time, etc.
    """

    def __init__(
        self,
        similarity_threshold: float = 0.95,
        default_ttl_seconds: Optional[int] = None,
        max_entries: int = 10000,
        encoder: Optional[Any] = None,
    ):
        """
        Initialize semantic cache.

        Args:
            similarity_threshold: Minimum similarity (0-1) to consider a cache hit
            default_ttl_seconds: Default TTL for cache entries (None = no expiry)
            max_entries: Maximum cache entries before eviction
            encoder: Optional encoder instance (creates default if None)
        """
        self.similarity_threshold = similarity_threshold
        self.default_ttl_seconds = default_ttl_seconds
        self.max_entries = max_entries

        # Initialize encoder
        if encoder is None:
            from hieromem_core.encoder import Encoder, EncoderConfig
            self._encoder = Encoder(EncoderConfig())
        else:
            self._encoder = encoder

        # Cache storage
        self._entries: Dict[str, CacheEntry] = {}
        self._embeddings: List[Tuple[str, np.ndarray]] = []  # (key, embedding)

        # Statistics
        self._stats = CacheStats()

    @property
    def stats(self) -> CacheStats:
        """Get cache statistics."""
        return self._stats

    def _compute_embedding(self, text: str) -> np.ndarray:
        """Compute embedding for text."""
        return self._encoder.encode(text)

    def _find_similar(
        self,
        embedding: np.ndarray,
        threshold: Optional[float] = None
    ) -> Optional[CacheEntry]:
        """
        Find a similar cached entry.

        Args:
            embedding: Query embedding
            threshold: Similarity threshold (uses default if None)

        Returns:
            CacheEntry if similar entry found, None otherwise
        """
        threshold = threshold or self.similarity_threshold

        if not self._embeddings:
            return None

        # Compute similarities
        best_key = None
        best_sim = -1.0

        for key, cached_emb in self._embeddings:
            # Cosine similarity
            sim = np.dot(embedding, cached_emb) / (
                np.linalg.norm(embedding) * np.linalg.norm(cached_emb) + 1e-8
            )
            if sim > best_sim:
                best_sim = sim
                best_key = key

        if best_sim >= threshold and best_key is not None:
            entry = self._entries.get(best_key)
            if entry and not entry.is_expired:
                return entry

        return None

    def _cache_key(self, prompt: str) -> str:
        """Generate cache key for prompt."""
        return hashlib.md5(prompt.encode(), usedforsecurity=False).hexdigest()

    def _evict_if_needed(self) -> None:
        """Evict old entries if cache is full."""
        if len(self._entries) < self.max_entries:
            return

        # Remove expired entries first
        expired_keys = [
            k for k, v in self._entries.items() if v.is_expired
        ]
        for key in expired_keys:
            self._remove_entry(key)

        # If still over limit, remove least recently used
        if len(self._entries) >= self.max_entries:
            # Sort by hits (ascending) and created_at (oldest first)
            sorted_entries = sorted(
                self._entries.items(),
                key=lambda x: (x[1].hits, x[1].created_at)
            )
            # Remove bottom 10%
            to_remove = len(self._entries) // 10
            for key, _ in sorted_entries[:to_remove]:
                self._remove_entry(key)

    def _remove_entry(self, key: str) -> None:
        """Remove a cache entry."""
        if key in self._entries:
            del self._entries[key]
        self._embeddings = [(k, e) for k, e in self._embeddings if k != key]

    def get(
        self,
        prompt: str,
        threshold: Optional[float] = None
    ) -> Optional[str]:
        """
        Get cached response for prompt.

        Args:
            prompt: The prompt to look up
            threshold: Optional similarity threshold override

        Returns:
            Cached response if found, None otherwise
        """
        embedding = self._compute_embedding(prompt)
        entry = self._find_similar(embedding, threshold)

        if entry:
            entry.hits += 1
            return entry.response

        return None

    def put(
        self,
        prompt: str,
        response: str,
        ttl_seconds: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Store a response in cache.

        Args:
            prompt: The prompt
            response: The LLM response
            ttl_seconds: Time-to-live (uses default if None)
            metadata: Optional metadata
        """
        self._evict_if_needed()

        embedding = self._compute_embedding(prompt)
        key = self._cache_key(prompt)

        entry = CacheEntry(
            prompt=prompt,
            response=response,
            embedding=embedding,
            created_at=datetime.now(),
            ttl_seconds=ttl_seconds or self.default_ttl_seconds,
            metadata=metadata or {}
        )

        self._entries[key] = entry
        self._embeddings.append((key, embedding))

    def get_or_compute(
        self,
        prompt: str,
        compute_fn: Callable[[str], str],
        threshold: Optional[float] = None,
        ttl_seconds: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> str:
        """
        Get cached response or compute and cache new one.

        This is the primary API for semantic caching.

        Args:
            prompt: The prompt
            compute_fn: Function to call if cache miss (prompt -> response)
            threshold: Optional similarity threshold override
            ttl_seconds: TTL for new cache entry
            metadata: Optional metadata for new entry

        Returns:
            The response (from cache or newly computed)
        """
        self._stats.total_requests += 1

        # Try cache lookup
        embedding = self._compute_embedding(prompt)
        entry = self._find_similar(embedding, threshold)

        if entry:
            self._stats.cache_hits += 1
            entry.hits += 1
            return entry.response

        # Cache miss - compute
        self._stats.cache_misses += 1
        start_time = time.perf_counter()
        response = compute_fn(prompt)
        elapsed_ms = (time.perf_counter() - start_time) * 1000

        # Store in cache
        self.put(prompt, response, ttl_seconds=ttl_seconds, metadata=metadata)

        # Track time that would be saved on future hits
        self._stats.total_saved_ms += elapsed_ms

        return response

    def clear(self) -> int:
        """
        Clear all cache entries.

        Returns:
            Number of entries cleared
        """
        count = len(self._entries)
        self._entries.clear()
        self._embeddings.clear()
        return count

    def prune_expired(self) -> int:
        """
        Remove expired entries.

        Returns:
            Number of entries removed
        """
        expired_keys = [
            k for k, v in self._entries.items() if v.is_expired
        ]
        for key in expired_keys:
            self._remove_entry(key)
        return len(expired_keys)

    def __len__(self) -> int:
        """Return number of cached entries."""
        return len(self._entries)
