# SPDX-License-Identifier: Apache-2.0
"""
MLLM (Multimodal Language Model) Prefix Cache Manager.

This module provides advanced caching for MLLM inference, implementing
the LMCache-style approach for multimodal prefix caching:

Features:
- Image content hashing for cache keys (LMCache style)
- Vision embedding caching (skip encoder on hit)
- KV cache state caching with prefix matching
- Token ID tracking for partial prefix reuse
- LRU eviction policy with memory limits
- Stats tracking (hits, misses, tokens saved, encoder skips)

Based on research from:
- LMCache: https://blog.lmcache.ai/2025-07-03-multimodal-models/
- vLLM Prefix Caching: https://docs.vllm.ai/en/stable/design/prefix_caching/
- mlx-lm cache_prompt: https://github.com/ml-explore/mlx-lm
"""

import hashlib
import logging
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class MLLMCacheStats:
    """Statistics for MLLM cache performance."""

    hits: int = 0
    misses: int = 0
    partial_hits: int = 0  # Prefix matched but not full
    tokens_saved: int = 0
    image_cache_hits: int = 0
    vision_encoder_skips: int = 0  # Times we skipped vision encoder
    total_queries: int = 0
    evictions: int = 0

    @property
    def hit_rate(self) -> float:
        """Calculate cache hit rate."""
        if self.total_queries == 0:
            return 0.0
        return self.hits / self.total_queries

    def to_dict(self) -> dict:
        """Convert stats to dictionary."""
        return {
            "hits": self.hits,
            "misses": self.misses,
            "partial_hits": self.partial_hits,
            "hit_rate": self.hit_rate,
            "tokens_saved": self.tokens_saved,
            "image_cache_hits": self.image_cache_hits,
            "vision_encoder_skips": self.vision_encoder_skips,
            "total_queries": self.total_queries,
            "evictions": self.evictions,
        }


@dataclass
class MLLMPrefixCacheEntry:
    """
    Enhanced cache entry storing vision embeddings, KV cache, and token IDs.

    This enables:
    1. Skipping vision encoder on image cache hit (saves ~1-2s per image)
    2. Skipping prefix computation on token match (saves ~0.5s per 1k tokens)
    3. Partial prefix reuse for multi-turn conversations
    """

    # Identity
    image_hash: str  # SHA256 of image content
    prompt_hash: str  # SHA256 of formatted prompt

    # Cached states - the key to performance
    vision_embeddings: Any = None  # Output of vision encoder (skip encoder on hit!)
    kv_cache: list[Any] = field(default_factory=list)  # Language model KV states

    # Token tracking for prefix matching
    token_ids: list[int] = field(default_factory=list)  # Full token sequence
    num_image_tokens: int = 0  # e.g., 256 for Gemma 3
    num_text_tokens: int = 0
    prompt_tokens: int = 0  # Total prompt tokens (for stats)

    # Metadata
    created_at: float = field(default_factory=time.time)
    hit_count: int = 0
    model_name: str = ""

    @property
    def total_tokens(self) -> int:
        return len(self.token_ids)

    @property
    def memory_size(self) -> int:
        """Estimate memory usage in bytes."""
        size = 0
        if self.vision_embeddings is not None:
            if hasattr(self.vision_embeddings, "nbytes"):
                size += self.vision_embeddings.nbytes
        if self.kv_cache is not None:
            for layer_cache in self.kv_cache:
                if hasattr(layer_cache, "state"):
                    state = layer_cache.state
                    if state is not None:
                        for tensor in state:
                            if hasattr(tensor, "nbytes"):
                                size += tensor.nbytes
        return size

    def get_prefix_match_length(self, new_token_ids: list[int]) -> int:
        """
        Find how many tokens match between cached prefix and new input.

        This is the key to prefix caching - if the first N tokens match,
        we can skip computing KV states for those N tokens.
        """
        match_length = 0
        for i, (cached, new) in enumerate(zip(self.token_ids, new_token_ids)):
            if cached != new:
                break
            match_length = i + 1
        return match_length


def compute_image_hash(image_path: str) -> str:
    """
    Compute hash of image content for cache key.

    Following LMCache approach: hash the actual image bytes, not the path.
    This ensures cache hits even when the same image is loaded from
    different paths or as base64.

    Args:
        image_path: Path to image file

    Returns:
        SHA256 hash of image content (first 16 chars)
    """
    try:
        path = Path(image_path)
        if path.exists():
            # Hash file content - this is the LMCache approach
            content = path.read_bytes()
            return hashlib.sha256(content).hexdigest()[:16]
        else:
            # Hash the string itself (for URLs or base64)
            return hashlib.sha256(image_path.encode()).hexdigest()[:16]
    except Exception as e:
        logger.warning(f"Failed to hash image: {e}")
        return hashlib.sha256(str(image_path).encode()).hexdigest()[:16]


def compute_images_hash(images: list[str]) -> str:
    """
    Compute combined hash for multiple images.

    Args:
        images: List of image paths/URLs

    Returns:
        Combined hash string
    """
    if not images:
        return "no_images"

    hashes = [compute_image_hash(img) for img in images]
    combined = "_".join(sorted(hashes))
    return hashlib.sha256(combined.encode()).hexdigest()[:16]


class MLLMPrefixCacheManager:
    """
    LRU Cache manager for MLLM prefix states with vision embedding caching.

    Implements the LMCache approach for multimodal caching:
    1. Hash-based identification of image+prompt combinations
    2. Vision embedding caching (skip encoder on hit - saves 1-2s!)
    3. KV cache reuse for matching prefixes
    4. Token ID tracking for partial prefix matching
    5. Memory-based eviction (configurable limit)

    Example:
        >>> cache = MLLMPrefixCacheManager(max_memory_mb=2048)
        >>> # First request - cache miss, full computation
        >>> entry, match_len = cache.fetch(["image.jpg"], prompt, token_ids)
        >>> # ... run full forward pass ...
        >>> cache.store(["image.jpg"], prompt, vision_emb, kv_cache, token_ids)
        >>>
        >>> # Second request with same image - cache hit!
        >>> entry, match_len = cache.fetch(["image.jpg"], prompt, token_ids)
        >>> # entry.vision_embeddings available - skip encoder!
        >>> # match_len > 0 - skip prefix computation!

    Performance (Gemma 3 27B, 256 image tokens):
        - Vision encoder: ~1.5s -> 0s (skip on hit)
        - Prefix computation: ~0.5s/1k tokens -> 0s (skip on match)
        - Multi-turn speedup: 8-12x for subsequent turns
    """

    def __init__(
        self,
        max_entries: int = 50,
        max_memory_mb: int = 2048,
    ):
        """
        Initialize MLLM prefix cache manager.

        Args:
            max_entries: Maximum number of cache entries (default: 50)
            max_memory_mb: Maximum memory in MB (default: 2048)
        """
        self.max_size = max_entries
        self.max_memory = max_memory_mb * 1024 * 1024
        self._cache: OrderedDict[str, MLLMPrefixCacheEntry] = OrderedDict()
        self._current_memory = 0
        self.stats = MLLMCacheStats()

    def _make_cache_key(self, images: list[str], prompt: str) -> str:
        """Create cache key from images and prompt."""
        image_hash = compute_images_hash(images)
        prompt_hash = hashlib.sha256(prompt.encode()).hexdigest()[:16]
        return f"{image_hash}_{prompt_hash}"

    def _make_image_only_key(self, images: list[str]) -> str:
        """Create cache key for image-only lookup (vision embedding reuse)."""
        return compute_images_hash(images)

    def _evict_by_memory(self, required_size: int) -> None:
        """Evict entries until we have enough memory."""
        while self._current_memory + required_size > self.max_memory and self._cache:
            oldest_key = next(iter(self._cache))
            oldest_entry = self._cache.pop(oldest_key)
            self._current_memory -= oldest_entry.memory_size
            self.stats.evictions += 1
            logger.debug(f"MLLM cache evicted (memory): {oldest_key[:20]}...")

    def _evict_by_count(self) -> None:
        """Evict entries until we're under max_size."""
        while len(self._cache) >= self.max_size and self._cache:
            oldest_key = next(iter(self._cache))
            oldest_entry = self._cache.pop(oldest_key)
            self._current_memory -= oldest_entry.memory_size
            self.stats.evictions += 1
            logger.debug(f"MLLM cache evicted (count): {oldest_key[:20]}...")

    def fetch(
        self,
        images: list[str],
        prompt: str,
        token_ids: list[int] | None = None,
    ) -> tuple[MLLMPrefixCacheEntry | None, int]:
        """
        Fetch cached prefix state with prefix matching.

        This is the main entry point for cache lookups. Returns both
        the cache entry (if found) and the prefix match length.

        Args:
            images: List of image paths
            prompt: Text prompt
            token_ids: Optional token IDs for prefix matching

        Returns:
            Tuple of (entry, prefix_match_length) where:
            - entry: The cache entry if found, None otherwise
            - prefix_match_length: Number of tokens that match (0 if miss)
        """
        self.stats.total_queries += 1
        cache_key = self._make_cache_key(images, prompt)

        if cache_key in self._cache:
            # Full cache hit - exact image+prompt match
            entry = self._cache.pop(cache_key)
            self._cache[cache_key] = entry  # Move to end (LRU)
            entry.hit_count += 1

            self.stats.hits += 1
            if images:
                self.stats.image_cache_hits += 1
            if entry.vision_embeddings is not None:
                self.stats.vision_encoder_skips += 1

            # Calculate prefix match length
            match_length = entry.total_tokens
            if token_ids:
                match_length = entry.get_prefix_match_length(token_ids)
                if match_length < entry.total_tokens:
                    self.stats.partial_hits += 1

            self.stats.tokens_saved += match_length
            logger.debug(
                f"MLLM cache HIT: {cache_key[:32]}..., prefix_match={match_length}"
            )

            return entry, match_length

        # Check for image-only match (can reuse vision embeddings)
        if images:
            image_key = self._make_image_only_key(images)
            for key, entry in self._cache.items():
                if (
                    entry.image_hash == image_key
                    and entry.vision_embeddings is not None
                ):
                    # Image match - can reuse vision embeddings!
                    self.stats.partial_hits += 1
                    self.stats.vision_encoder_skips += 1
                    logger.debug(
                        f"MLLM cache PARTIAL HIT (vision only): image={image_key[:16]}"
                    )

                    # Return entry for vision embeddings, but 0 prefix match
                    # (prompt is different, so KV cache can't be reused)
                    return entry, 0

        self.stats.misses += 1
        logger.debug(f"MLLM cache MISS: {cache_key[:32]}...")
        return None, 0

    def fetch_cache(
        self,
        images: list[str],
        prompt: str,
    ) -> tuple[list[Any] | None, bool]:
        """
        Legacy API: Fetch cached KV state for image+prompt combination.

        For backwards compatibility with existing code.
        """
        entry, match_len = self.fetch(images, prompt)
        # For legacy API, return hit if entry exists (don't require token match)
        if entry is not None and entry.kv_cache is not None:
            return entry.kv_cache, True
        return None, False

    def store(
        self,
        images: list[str],
        prompt: str,
        vision_embeddings: Any,
        kv_cache: list[Any],
        token_ids: list[int],
        num_image_tokens: int = 0,
        model_name: str = "",
    ) -> None:
        """
        Store prefix state in cache.

        Args:
            images: List of image paths
            prompt: Text prompt
            vision_embeddings: Output of vision encoder (can be None for text-only)
            kv_cache: Language model KV cache states
            token_ids: Full token sequence
            num_image_tokens: Number of image tokens (e.g., 256 for Gemma 3)
            model_name: Model name for validation
        """
        cache_key = self._make_cache_key(images, prompt)

        entry = MLLMPrefixCacheEntry(
            image_hash=compute_images_hash(images),
            prompt_hash=hashlib.sha256(prompt.encode()).hexdigest()[:16],
            vision_embeddings=vision_embeddings,
            kv_cache=kv_cache,
            token_ids=token_ids,
            num_image_tokens=num_image_tokens,
            num_text_tokens=len(token_ids) - num_image_tokens,
            prompt_tokens=len(token_ids),
            model_name=model_name,
        )

        # Evict by memory first
        self._evict_by_memory(entry.memory_size)

        # Then evict by count
        self._evict_by_count()

        self._cache[cache_key] = entry
        self._current_memory += entry.memory_size

        logger.debug(
            f"MLLM cache STORED: key={cache_key[:32]}..., "
            f"tokens={len(token_ids)}, vision_emb={vision_embeddings is not None}, "
            f"memory={entry.memory_size / 1024 / 1024:.1f}MB"
        )

    def store_cache(
        self,
        images: list[str],
        prompt: str,
        cache: list[Any] | None,
        num_tokens: int = 0,
    ) -> None:
        """
        Legacy API: Store KV cache for future reuse.

        For backwards compatibility with existing code.
        """
        # Don't store empty or None caches
        if cache is None or (isinstance(cache, list) and len(cache) == 0):
            return

        self.store(
            images=images,
            prompt=prompt,
            vision_embeddings=None,
            kv_cache=cache,
            token_ids=[0] * num_tokens,  # Dummy token IDs
            num_image_tokens=0,
        )

    def get_stats(self) -> dict[str, Any]:
        """Get cache statistics."""
        stats = self.stats.to_dict()
        stats["entries"] = len(self._cache)
        stats["max_entries"] = self.max_size
        stats["memory_used_mb"] = self._current_memory / 1024 / 1024
        stats["max_memory_mb"] = self.max_memory / 1024 / 1024
        return stats

    def reset_stats(self) -> None:
        """Reset statistics counters."""
        self.stats = MLLMCacheStats()

    def clear(self) -> None:
        """Clear all cached entries and reset stats."""
        self._cache.clear()
        self._current_memory = 0
        self.reset_stats()

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

    def __repr__(self) -> str:
        mem_mb = self._current_memory / 1024 / 1024
        return f"<MLLMPrefixCacheManager entries={len(self)} memory={mem_mb:.1f}MB>"


# Short alias for convenience
MLLMCacheManager = MLLMPrefixCacheManager

# Legacy aliases for backwards compatibility
VLMCacheStats = MLLMCacheStats
VLMPrefixCacheEntry = MLLMPrefixCacheEntry
VLMCacheEntry = MLLMPrefixCacheEntry
VLMPrefixCacheManager = MLLMPrefixCacheManager
VLMCacheManager = MLLMPrefixCacheManager
