"""
Cache utilities.
"""

import pickle

import cachetools
import zstandard as zstd

# Singleton instances of zstandard compressors and decompressors
_ZSTD_COMPRESSOR = zstd.ZstdCompressor()
_ZSTD_DECOMPRESSOR = zstd.ZstdDecompressor()


class TTLCache(cachetools.TTLCache):
    """A TTLCache that supports compression using zstandard and pickle."""

    def __init__(self, maxsize: int | float, ttl: int | float, compress: bool = False):
        super().__init__(maxsize=float(maxsize), ttl=float(ttl))
        self._compress_enabled = compress

    def __setitem__(self, key, value):
        if self._compress_enabled:
            value = _ZSTD_COMPRESSOR.compress(
                pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
            )
        super().__setitem__(key, value)

    def __getitem__(self, key):
        value = super().__getitem__(key)
        if self._compress_enabled:
            value = pickle.loads(_ZSTD_DECOMPRESSOR.decompress(value))
        return value