content: #  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Memory cache module for the ZenML server."""
  
  import time
  from collections import OrderedDict
  from threading import Lock
  from typing import Any, Callable, Optional
  from typing import OrderedDict as OrderedDictType
  from uuid import UUID
  
  from zenml.logger import get_logger
  from zenml.utils.singleton import SingletonMetaClass
  
  logger = get_logger(__name__)
  
  
  class MemoryCacheEntry:
      """Simple class to hold cache entry data."""
  
      def __init__(self, value: Any, expiry: int) -> None:
          """Initialize a cache entry with value and expiry time.
  
          Args:
              value: The value to store in the cache.
              expiry: The expiry time in seconds.
          """
          self.value: Any = value
          self.expiry: int = expiry
          self.timestamp: float = time.time()
  
      @property
      def expired(self) -> bool:
          """Check if the cache entry has expired.
  
          Returns:
              True if the cache entry has expired; otherwise, False.
          """
          return time.time() - self.timestamp >= self.expiry
  
  
  class MemoryCache(metaclass=SingletonMetaClass):
      """Simple in-memory cache with expiry and capacity management.
  
      This cache is thread-safe and can be used in both synchronous and
      asynchronous contexts. It uses a simple LRU (Least Recently Used) eviction
      strategy to manage the cache size.
  
      Each cache entry has a key, value, timestamp, and expiry. The cache
      automatically removes expired entries and evicts the oldest entry when
      the cache reaches its maximum capacity.
  
  
      Usage Example:
  
          cache = MemoryCache()
          uuid_key = UUID("12345678123456781234567812345678")
  
          if not cache.get(uuid_key):
              # Get the value from the database or other source
              value = get_value_from_database()
              cache.set(uuid_key, value, expiry=60)
  
      Usage Example with decorator:
  
          @cache_result(expiry=60)
          def get_cached_value(key: UUID) -> Any:
              return get_value_from_database(key)
  
          uuid_key = UUID("12345678123456781234567812345678")
  
          value = get_cached_value(uuid_key)
      """
  
      def __init__(self, max_capacity: int, default_expiry: int) -> None:
          """Initialize the cache with a maximum capacity and default expiry time.
  
          Args:
              max_capacity: The maximum number of entries the cache can hold.
              default_expiry: The default expiry time in seconds.
          """
          self.cache: OrderedDictType[UUID, MemoryCacheEntry] = OrderedDict()
          self.max_capacity = max_capacity
          self.default_expiry = default_expiry
          self._lock = Lock()
  
      def set(self, key: UUID, value: Any, expiry: Optional[int] = None) -> None:
          """Insert value into cache with optional custom expiry time in seconds.
  
          Args:
              key: The key to insert the value with.
              value: The value to insert into the cache.
              expiry: The expiry time in seconds. If None, uses the default expiry.
          """
          with self._lock:
              self.cache[key] = MemoryCacheEntry(
                  value=value, expiry=expiry or self.default_expiry
              )
              self._cleanup()
  
      def get(self, key: UUID) -> Optional[Any]:
          """Retrieve value if it's still valid; otherwise, return None.
  
          Args:
              key: The key to retrieve the value for.
  
          Returns:
              The value if it's still valid; otherwise, None.
          """
          with self._lock:
              return self._get_internal(key)
  
      def _get_internal(self, key: UUID) -> Optional[Any]:
          """Helper to retrieve a value without lock (internal use only).
  
          Args:
              key: The key to retrieve the value for.
  
          Returns:
              The value if it's still valid; otherwise, None.
          """
          entry = self.cache.get(key)
          if entry and not entry.expired:
              return entry.value
          elif entry:
              del self.cache[key]  # Invalidate expired entry
          return None
  
      def _cleanup(self) -> None:
          """Remove expired or excess entries."""
          # Remove expired entries
          keys_to_remove = [k for k, v in self.cache.items() if v.expired]
          for k in keys_to_remove:
              del self.cache[k]
  
          # Ensure we don't exceed max capacity
          while len(self.cache) > self.max_capacity:
              self.cache.popitem(last=False)
  
  
  F = Callable[[UUID], Any]
  
  
  def cache_result(
      expiry: Optional[int] = None,
  ) -> Callable[[F], F]:
      """A decorator to cache the result of a function based on a UUID key argument.
  
      Args:
          expiry: Custom time in seconds for the cache entry to expire. If None,
              uses the default expiry time.
  
      Returns:
          A decorator that wraps a function, caching its results based on a UUID
          key.
      """
  
      def decorator(func: F) -> F:
          """The actual decorator that wraps the function with caching logic.
  
          Args:
              func: The function to wrap.
  
          Returns:
              The wrapped function with caching logic.
          """
  
          def wrapper(key: UUID) -> Any:
              """The wrapped function with caching logic.
  
              Args:
                  key: The key to use for caching.
  
              Returns:
                  The result of the original function, either from cache or
                  freshly computed.
              """
              from zenml.zen_server.utils import memcache
  
              cache = memcache()
  
              # Attempt to retrieve the result from cache
              cached_value = cache.get(key)
              if cached_value is not None:
                  logger.debug(
                      f"Memory cache hit for key: {key} and func: {func.__name__}"
                  )
                  return cached_value
  
              # Call the original function and cache its result
              result = func(key)
              cache.set(key, result, expiry)
              return result
  
          return wrapper
  
      return decorator
  