import hashlib
from typing import Dict, Optional


class PromptCache:
    """
    A class to manage prompt caching for LLMs, allowing the reuse of context across multiple API requests.

    This reduces costs and latency, particularly for long prompts.

    Attributes:
        cache (Dict[str, str]): A dictionary to store cached prompts and their corresponding responses.
    """

    def __init__(self) -> None:
        """Initializes the PromptCache with an empty cache."""
        self.cache: Dict[str, str] = {}

    def _hash_prompt(self, prompt: str) -> str:
        """
        Generates a unique hash for a given prompt.

        Args:
            prompt (str): The prompt to hash.

        Returns:
            str: The generated hash.
        """
        return hashlib.sha256(prompt.encode()).hexdigest()

    def add_to_cache(self, prompt: str, response: str) -> None:
        """
        Adds a prompt and its corresponding response to the cache.

        Args:
            prompt (str): The prompt string.
            response (str): The response generated by the LLM.

        Returns:
            None
        """
        prompt_hash = self._hash_prompt(prompt)
        self.cache[prompt_hash] = response

    def get_from_cache(self, prompt: str) -> Optional[str]:
        """
        Retrieves a cached response for a given prompt, if available.

        Args:
            prompt (str): The prompt string to retrieve the cached response for.

        Returns:
            Optional[str]: The cached response if found, otherwise None.
        """
        prompt_hash = self._hash_prompt(prompt)
        return self.cache.get(prompt_hash)

    def clear_cache(self) -> None:
        """
        Clears the entire prompt cache.

        Returns:
            None
        """
        self.cache.clear()

    def cache_size(self) -> int:
        """
        Returns the number of items currently in the cache.

        Returns:
            int: The size of the cache.
        """
        return len(self.cache)

    def remove_from_cache(self, prompt: str) -> None:
        """
        Removes a specific prompt and its response from the cache.

        Args:
            prompt (str): The prompt string to remove from the cache.

        Returns:
            None
        """
        prompt_hash = self._hash_prompt(prompt)
        if prompt_hash in self.cache:
            del self.cache[prompt_hash]


# Example usage:

# Initialize the cache
prompt_cache = PromptCache()

# Add a prompt and response to the cache
prompt = "What is the capital of France?"
response = "The capital of France is Paris."
prompt_cache.add_to_cache(prompt, response)

# Retrieve the response from the cache
cached_response = prompt_cache.get_from_cache(prompt)
if cached_response:
    print("Cached response:", cached_response)
else:
    print("Prompt not found in cache.")
