from typing import Optional, List, Dict, Any
from datetime import datetime, timedelta
from uuid import UUID, uuid4
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity as sk_cosine_similarity
from sentence_transformers import SentenceTransformer
from fastapi import BackgroundTasks
from inferencekit.core.types import (
    Request,
    Response,
    CacheEntry,
    CacheConfig,
    SimilarityResult,
    ModelType,
    ProviderType,
    ProviderConfig,
    ModelConfig,
    BatchRequest,
    BatchResponse,
)
from inferencekit.core.exceptions import (
    CacheError,
    CacheFullError,
    CacheNotFoundError,
    CacheExpiredError,
    ModelNotFoundError,
    ProviderError,
    AuthenticationError,
    RateLimitError,
    InvalidRequestError,
    ConnectionError,
    OperationTimeoutError,
)
from inferencekit.backends.memory_backend import MemoryBackend
from inferencekit.providers.openai import OpenAIProvider




class SemanticCache:
    """Semantic cache implementation with vector similarity"""

    def __init__(self, config: CacheConfig):
        self.config = config
        self.entries: Dict[UUID, CacheEntry] = {}
        self.vector_store: List[float] = []
        self.transformer_model: Optional[SentenceTransformer] = None
        self.ttl_store: Dict[UUID, datetime] = {}

        if config.enabled:
            self._load_transformer_model()

    def _load_transformer_model(self) -> None:
        """Load transformer model for embeddings"""
        try:
            self.transformer_model = SentenceTransformer('all-MiniLM-L6-v2')
        except Exception as e:
            raise CacheError(f"Failed to load transformer model: {e}")

    def find_similar(self, request: Request) -> SimilarityResult:
        """Find similar cached entries"""
        if not self.transformer_model:
            return SimilarityResult(request_id=request.id, similar_entries=[], scores=[])

        try:
            # Generate embedding for the request
            embedding = self.transformer_model.encode(request.prompt)

            # Calculate similarity scores
            scores = []
            similar_entries = []

            for entry_id, entry in self.entries.items():
                if entry_id in self.ttl_store and self.ttl_store[entry_id] < datetime.utcnow():
                    continue

                # Calculate cosine similarity
                entry_embedding = entry.embedding
                if entry_embedding:
                    score = np.dot(embedding, entry_embedding) / (np.linalg.norm(embedding) * np.linalg.norm(entry_embedding))
                    scores.append(score)
                    similar_entries.append(entry)

            return SimilarityResult(request_id=request.id, similar_entries=similar_entries, scores=scores)
        except Exception as e:
            raise CacheError(f"Failed to find similar entries: {e}")

    def add_entry(self, request: Request, response: Response) -> None:
        """Add entry to cache"""
        if not self.transformer_model:
            return

        try:
            # Generate embedding
            embedding = self.transformer_model.encode(request.prompt)

            # Create cache entry
            entry = CacheEntry(
                request_id=request.id,
                request=request,
                response=response,
                embedding=embedding,
                created_at=datetime.utcnow()
            )

            # Add to cache
            self.entries[request.id] = entry
            self.ttl_store[request.id] = datetime.utcnow() + timedelta(seconds=self.config.ttl_seconds)

        except Exception as e:
            raise CacheError(f"Failed to add cache entry: {e}")

    def get_entry(self, request_id: UUID) -> Optional[CacheEntry]:
        """Get cache entry by request ID"""
        entry = self.entries.get(request_id)
        if entry and request_id in self.ttl_store:
            if self.ttl_store[request_id] > datetime.utcnow():
                return entry
            else:
                # Remove expired entry
                del self.entries[request_id]
                del self.ttl_store[request_id]
        return None

    def clear(self) -> None:
        """Clear all cache entries"""
        self.entries.clear()
        self.vector_store.clear()
        self.ttl_store.clear()

    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics"""
        return {
            "total_entries": len(self.entries),
            "expired_entries": sum(1 for ts in self.ttl_store.values() if ts < datetime.utcnow()),
            "memory_usage": "N/A"  # This would need implementation
        }

    def _cleanup_expired_entries(self) -> None:
        """Cleanup expired cache entries"""
        now = datetime.utcnow()
        expired_ids = [req_id for req_id, ts in self.ttl_store.items() if ts < now]

        for req_id in expired_ids:
            self.entries.pop(req_id, None)
            self.ttl_store.pop(req_id, None)

    def background_cleanup(self, background_tasks: BackgroundTasks) -> None:
        """Schedule background cleanup task"""
        background_tasks.add_task(self._cleanup_expired_entries())


class Cache:
    """Factory for creating cache instances"""

    @staticmethod
    def semantic(
        backend: str = "memory",
        similarity_threshold: float = 0.92,
        ttl_seconds: int = 3600,
        **kwargs
    ) -> SemanticCache:
        """
        Create semantic cache with sensible defaults.

        Args:
            backend: Cache backend ("redis", "sqlite", "memory")
            similarity_threshold: Minimum similarity score for cache hits
            ttl_seconds: Time-to-live for cache entries in seconds
            **kwargs: Additional configuration options
        """
        config = CacheConfig(
            enabled=True,
            backend=backend,
            similarity_threshold=similarity_threshold,
            ttl_seconds=ttl_seconds,
            **kwargs
        )
        return SemanticCache(config)
    """Semantic cache implementation with vector similarity"""

    def __init__(self, config: CacheConfig):
        self.config = config
        self.entries: Dict[UUID, CacheEntry] = {}
        self.vector_store: List[float] = []
        self.transformer_model: Optional[SentenceTransformer] = None
        self.ttl_store: Dict[UUID, datetime] = {}

        if config.enabled:
            self._load_transformer_model()

    def _load_transformer_model(self) -> None:
        """Load transformer model for embeddings"""
        try:
            self.transformer_model = SentenceTransformer('all-MiniLM-L6-v2')
        except Exception as e:
            raise CacheError(f"Failed to load transformer model: {e}")

    def find_similar(self, request: Request) -> SimilarityResult:
        """Find similar cached entries"""
        if not self.transformer_model:
            return SimilarityResult(request_id=request.id, similar_entries=[], scores=[])

        try:
            # Generate embedding for the request
            embedding = self.transformer_model.encode(request.prompt)

            # Calculate similarity scores
            scores = []
            similar_entries = []

            for entry_id, entry in self.entries.items():
                if entry_id in self.ttl_store and self.ttl_store[entry_id] < datetime.utcnow():
                    continue

                # Calculate cosine similarity
                entry_embedding = entry.embedding
                if entry_embedding:
                    score = np.dot(embedding, entry_embedding) / (np.linalg.norm(embedding) * np.linalg.norm(entry_embedding))
                    scores.append(score)
                    similar_entries.append(entry)

            return SimilarityResult(request_id=request.id, similar_entries=similar_entries, scores=scores)
        except Exception as e:
            raise CacheError(f"Failed to find similar entries: {e}")

    def add_entry(self, request: Request, response: Response) -> None:
        """Add entry to cache"""
        if not self.transformer_model:
            return

        try:
            # Generate embedding
            embedding = self.transformer_model.encode(request.prompt)

            # Create cache entry
            entry = CacheEntry(
                request_id=request.id,
                request=request,
                response=response,
                embedding=embedding,
                created_at=datetime.utcnow()
            )

            # Add to cache
            self.entries[request.id] = entry
            self.ttl_store[request.id] = datetime.utcnow() + timedelta(seconds=self.config.ttl_seconds)

        except Exception as e:
            raise CacheError(f"Failed to add cache entry: {e}")

    def get_entry(self, request_id: UUID) -> Optional[CacheEntry]:
        """Get cache entry by request ID"""
        entry = self.entries.get(request_id)
        if entry and request_id in self.ttl_store:
            if self.ttl_store[request_id] > datetime.utcnow():
                return entry
            else:
                # Remove expired entry
                del self.entries[request_id]
                del self.ttl_store[request_id]
        return None

    def clear(self) -> None:
        """Clear all cache entries"""
        self.entries.clear()
        self.vector_store.clear()
        self.ttl_store.clear()

    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics"""
        return {
            "total_entries": len(self.entries),
            "expired_entries": sum(1 for ts in self.ttl_store.values() if ts < datetime.utcnow()),
            "memory_usage": "N/A"  # This would need implementation
        }

    def _cleanup_expired_entries(self) -> None:
        """Cleanup expired cache entries"""
        now = datetime.utcnow()
        expired_ids = [req_id for req_id, ts in self.ttl_store.items() if ts < now]

        for req_id in expired_ids:
            self.entries.pop(req_id, None)
            self.ttl_store.pop(req_id, None)

    def background_cleanup(self, background_tasks: BackgroundTasks) -> None:
        """Schedule background cleanup task"""
        background_tasks.add_task(self._cleanup_expired_entries())