"""Caching and batch processing for embeddings."""
import logging
import hashlib
import numpy as np
from typing import Dict, List, Tuple, Optional
import json
from pathlib import Path

logger = logging.getLogger(__name__)


class EmbeddingCache:
    """
    Cache embeddings to avoid recomputing.

    Stores hash(text) → embedding mapping
    """

    def __init__(self, cache_dir: Optional[str] = None):
        """Initialize cache."""
        self.cache_dir = Path(cache_dir) if cache_dir else None
        self.memory_cache: Dict[str, np.ndarray] = {}
        self.hits = 0
        self.misses = 0

        if self.cache_dir:
            self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _hash_text(self, text: str) -> str:
        """Generate hash for text."""
        return hashlib.md5(text.encode()).hexdigest()

    def get(self, text: str) -> Optional[np.ndarray]:
        """Get cached embedding for text."""
        text_hash = self._hash_text(text)

        # Try memory cache first
        if text_hash in self.memory_cache:
            self.hits += 1
            return self.memory_cache[text_hash]

        # Try disk cache
        if self.cache_dir:
            cache_file = self.cache_dir / f"{text_hash}.npy"
            if cache_file.exists():
                try:
                    embedding = np.load(cache_file)
                    self.memory_cache[text_hash] = embedding
                    self.hits += 1
                    return embedding
                except Exception as e:
                    logger.debug(f"Error loading cache {cache_file}: {e}")

        self.misses += 1
        return None

    def set(self, text: str, embedding: np.ndarray) -> bool:
        """Cache embedding for text."""
        text_hash = self._hash_text(text)
        self.memory_cache[text_hash] = embedding

        # Optionally save to disk
        if self.cache_dir:
            try:
                cache_file = self.cache_dir / f"{text_hash}.npy"
                np.save(cache_file, embedding)
            except Exception as e:
                logger.debug(f"Error saving cache: {e}")
                return False

        return True

    def get_stats(self) -> Dict[str, any]:
        """Get cache statistics."""
        total_lookups = self.hits + self.misses
        hit_rate = (self.hits / total_lookups * 100) if total_lookups > 0 else 0

        return {
            'hits': self.hits,
            'misses': self.misses,
            'hit_rate': hit_rate,
            'memory_cache_size': len(self.memory_cache),
            'total_lookups': total_lookups
        }

    def clear(self) -> None:
        """Clear memory cache."""
        self.memory_cache.clear()
        self.hits = 0
        self.misses = 0


class EmbeddingBatcher:
    """
    Batch process embeddings efficiently.

    Processes multiple emails' embeddings in parallel.
    """

    def __init__(self, embedder, batch_size: int = 32):
        """Initialize batcher."""
        self.embedder = embedder
        self.batch_size = batch_size
        self.cache = EmbeddingCache()

    def batch_encode(self, texts: List[str]) -> np.ndarray:
        """
        Encode texts in batches.

        Returns:
            Array of shape (len(texts), embedding_dim)
        """
        if not texts:
            return np.array([])

        embeddings = []

        # Check cache first
        uncached_indices = []
        uncached_texts = []

        for i, text in enumerate(texts):
            cached = self.cache.get(text)
            if cached is not None:
                embeddings.append((i, cached))
            else:
                uncached_indices.append(i)
                uncached_texts.append(text)

        # Process uncached in batches
        if uncached_texts:
            logger.debug(f"Processing {len(uncached_texts)} uncached embeddings in batches of {self.batch_size}")

            for batch_start in range(0, len(uncached_texts), self.batch_size):
                batch_end = min(batch_start + self.batch_size, len(uncached_texts))
                batch_texts = uncached_texts[batch_start:batch_end]
                batch_indices = uncached_indices[batch_start:batch_end]

                try:
                    batch_embeddings = self.embedder.encode(
                        batch_texts,
                        convert_to_numpy=True,
                        batch_size=self.batch_size
                    )

                    # Cache and collect
                    for idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
                        self.cache.set(text, embedding)
                        embeddings.append((idx, embedding))

                except Exception as e:
                    logger.error(f"Error encoding batch: {e}")
                    # Return zeros for failed embeddings
                    for idx in batch_indices:
                        embeddings.append((idx, np.zeros(384)))

        # Sort by original index
        embeddings.sort(key=lambda x: x[0])
        result = np.array([e[1] for e in embeddings])

        logger.debug(f"Cache stats: {self.cache.get_stats()}")

        return result

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