from __future__ import annotations

import logging
from collections import defaultdict
from collections.abc import (
    Callable,
    Generator,
    Hashable,
    Mapping,
    MutableMapping,
    Sized,
)
from contextlib import contextmanager
from functools import partial
from typing import Literal, NamedTuple, Protocol, cast

import zict
from dask.typing import Key

from distributed.metrics import context_meter
from distributed.protocol import deserialize_bytes, serialize_bytelist
from distributed.protocol.compression import get_compression_settings
from distributed.sizeof import safe_sizeof
from distributed.utils import RateLimiterFilter, nbytes

logger = logging.getLogger(__name__)
logger.addFilter(RateLimiterFilter("Spill file on disk reached capacity"))
logger.addFilter(RateLimiterFilter("Spill to disk failed"))


class SpilledSize(NamedTuple):
    """Size of a key/value pair when spilled to disk, in bytes"""

    # output of sizeof()
    memory: int
    # pickled size
    disk: int

    def __add__(self, other: SpilledSize) -> SpilledSize:  # type: ignore
        return SpilledSize(self.memory + other.memory, self.disk + other.disk)

    def __sub__(self, other: SpilledSize) -> SpilledSize:
        return SpilledSize(self.memory - other.memory, self.disk - other.disk)


class ManualEvictProto(Protocol):
    """Duck-type API that a third-party alternative to SpillBuffer must respect (in
    addition to MutableMapping) if it wishes to support spilling when the
    ``distributed.worker.memory.spill`` threshold is surpassed.

    This is public API. At the moment of writing, Dask-CUDA implements this protocol in
    the ProxifyHostFile class.
    """

    @property
    def fast(self) -> Sized | bool:
        """Access to fast memory. This is normally a MutableMapping, but for the purpose
        of the manual eviction API it is just tested for emptiness to know if there is
        anything to evict.
        """
        ...  # pragma: nocover

    def evict(self) -> int:
        """Manually evict a key/value pair from fast to slow memory.
        Return size of the evicted value in fast memory.

        If the eviction failed for whatever reason, return -1. This method must
        guarantee that the key/value pair that caused the issue has been retained in
        fast memory and that the problem has been logged internally.

        This method never raises.
        """
        ...  # pragma: nocover


class SpillBuffer(zict.Buffer[Key, object]):
    """MutableMapping that automatically spills out dask key/value pairs to disk when
    the total size of the stored data exceeds the target. If max_spill is provided the
    key/value pairs won't be spilled once this threshold has been reached.

    Parameters
    ----------
    spill_directory: str
        Location on disk to write the spill files to
    target: int
        Managed memory, in bytes, to start spilling at
    max_spill: int | False, optional
        Limit of number of bytes to be spilled on disk. Set to False to disable.
    """

    logged_pickle_errors: set[Key]
    #: (label, unit) -> ever-increasing cumulative value
    cumulative_metrics: defaultdict[tuple[str, str], float]

    def __init__(
        self,
        spill_directory: str,
        target: int,
        max_spill: int | Literal[False] = False,
    ):
        # If a value is still in use somewhere on the worker since the last time it was
        # unspilled, don't duplicate it
        slow = Slow(spill_directory, max_spill)
        slow_cached = zict.Cache(slow, zict.WeakValueMapping())

        super().__init__(fast={}, slow=slow_cached, n=target, weight=_in_memory_weight)
        self.logged_pickle_errors = set()  # keys logged with pickle error
        self.cumulative_metrics = defaultdict(float)

    @contextmanager
    def _capture_metrics(self) -> Generator[None]:
        """Capture metrics re. disk read/write, serialize/deserialize, and
        compress/decompress.

        Note that this duplicates capturing from gather_dep, get_data, and execute. It
        is repeated here to make it possible to split serialize/deserialize and
        compress/decompress triggered by spill/unspill from those triggered by network
        comms.
        """

        def metrics_callback(label: Hashable, value: float, unit: str) -> None:
            assert isinstance(label, str)
            self.cumulative_metrics[label, unit] += value

        with context_meter.add_callback(metrics_callback):
            yield

    @contextmanager
    def _handle_errors(self, key: Key | None) -> Generator[None]:
        try:
            yield
        except MaxSpillExceeded as e:
            # key is in self.fast; no keys have been lost on eviction
            (key_e,) = e.args
            assert key_e in self.fast
            assert key_e not in self.slow
            logger.warning(
                "Spill file on disk reached capacity; keeping data in memory"
            )
            raise HandledError()
        except OSError:
            # Typically, this is a disk full error
            logger.error("Spill to disk failed; keeping data in memory", exc_info=True)
            raise HandledError()
        except PickleError as e:
            assert e.key in self.fast
            assert e.key not in self.slow
            if e.key == key:
                assert key is not None
                # The key we just inserted failed to serialize.
                # This happens only when the key is individually larger than target.
                # The exception will be caught by Worker and logged; the status of
                # the task will be set to error.
                del self[key]
                raise
            else:
                # The key we just inserted is smaller than target, but it caused
                # another, unrelated key to be spilled out of the LRU, and that key
                # failed to serialize. There's nothing wrong with the new key. The older
                # key is still in memory.
                if e.key not in self.logged_pickle_errors:
                    logger.error("Failed to pickle %r", e.key, exc_info=True)
                    self.logged_pickle_errors.add(e.key)
                raise HandledError()

    def __setitem__(self, key: Key, value: object) -> None:
        """If sizeof(value) < target, write key/value pair to self.fast; this may in
        turn cause older keys to be spilled from fast to slow.
        If sizeof(value) >= target, write key/value pair directly to self.slow instead.

        Raises
        ------
        Exception
            sizeof(value) >= target, and value failed to pickle.
            The key/value pair has been forgotten.

        In all other cases:

        - an older value was evicted and failed to pickle,
        - this value or an older one caused the disk to fill and raise OSError,
        - this value or an older one caused the max_spill threshold to be exceeded,

        this method does not raise and guarantees that the key/value that caused the
        issue remained in fast.
        """
        try:
            with self._capture_metrics(), self._handle_errors(key):
                super().__setitem__(key, value)
                self.logged_pickle_errors.discard(key)
        except HandledError:
            assert key in self.fast
            assert key not in self.slow

    def evict(self) -> int:
        """Implementation of :meth:`ManualEvictProto.evict`.

        Manually evict the oldest key/value pair, even if target has not been
        reached. Returns sizeof(value).
        If the eviction failed (value failed to pickle, disk full, or max_spill
        exceeded), return -1; the key/value pair that caused the issue will remain in
        fast. The exception has been logged internally.
        This method never raises.
        """
        try:
            with self._capture_metrics(), self._handle_errors(None):
                _, _, weight = self.fast.evict()
                return cast(int, weight)
        except HandledError:
            return -1

    def __getitem__(self, key: Key) -> object:
        with self._capture_metrics():
            if key in self.fast:
                # Note: don't log from self.fast.__getitem__, because that's called
                # every time a key is evicted, and we don't want to count those events
                # here.
                memory_size = cast(int, self.fast.weights[key])
                # This is logged not only by the internal metrics callback but also by
                # those installed by gather_dep, get_data, and execute
                context_meter.digest_metric("memory-read", 1, "count")
                context_meter.digest_metric("memory-read", memory_size, "bytes")

            return super().__getitem__(key)

    def __delitem__(self, key: Key) -> None:
        super().__delitem__(key)
        self.logged_pickle_errors.discard(key)

    def pop(self, key: Key, default: object = None) -> object:
        raise NotImplementedError(
            "Are you calling .pop(key, None) as a way to discard a key if it exists?"
            "It may cause data to be read back from disk! Please use `del` instead."
        )

    @property
    def memory(self) -> Mapping[Key, object]:
        """Key/value pairs stored in RAM. Alias of zict.Buffer.fast.
        For inspection only - do not modify directly!
        """
        return self.fast

    @property
    def disk(self) -> Mapping[Key, object]:
        """Key/value pairs spilled out to disk. Alias of zict.Buffer.slow.
        For inspection only - do not modify directly!
        """
        return self.slow

    @property
    def _slow_uncached(self) -> Slow:
        cache = cast(zict.Cache, self.slow)
        return cast(Slow, cache.data)

    @property
    def spilled_total(self) -> SpilledSize:
        """Number of bytes spilled to disk. Tuple of

        - output of sizeof()
        - pickled size

        The two may differ substantially, e.g. if sizeof() is inaccurate or in case of
        compression.
        """
        return self._slow_uncached.total_weight


def _in_memory_weight(key: Key, value: object) -> int:
    return safe_sizeof(value)


# Internal exceptions. These are never raised by SpillBuffer.
class MaxSpillExceeded(Exception):
    pass


class PickleError(TypeError):
    def __str__(self) -> str:
        return f"Failed to pickle {self.key!r}"

    @property
    def key(self) -> Key:
        return self.args[0]


class HandledError(Exception):
    pass


class AnyKeyFile(zict.File):
    def _safe_key(self, key: Key) -> str:
        # We don't need _proper_ stringification, just a unique mapping
        return super()._safe_key(str(key))


class Slow(zict.Func[Key, object, bytes]):
    max_weight: int | Literal[False]
    weight_by_key: dict[Key, SpilledSize]
    total_weight: SpilledSize

    def __init__(self, spill_directory: str, max_weight: int | Literal[False] = False):
        compression = get_compression_settings(
            "distributed.worker.memory.spill-compression"
        )

        # File is MutableMapping[str, bytes], but serialize_bytelist returns
        # list[bytes | bytearray | memorymapping], which File.__setitem__ actually
        # accepts despite its signature; File.__getitem__ actually returns
        # bytearray. This headache is because MutableMapping doesn't allow for
        # asymmetric VT in __getitem__ and __setitem__.
        dump = cast(
            Callable[[object], bytes],
            partial(serialize_bytelist, compression=compression, on_error="raise"),
        )
        super().__init__(
            dump,
            deserialize_bytes,
            cast(MutableMapping[Key, bytes], AnyKeyFile(spill_directory)),
        )
        self.max_weight = max_weight
        self.weight_by_key = {}
        self.total_weight = SpilledSize(0, 0)

    def __getitem__(self, key: Key) -> object:
        with context_meter.meter("disk-read", "seconds"):
            pickled = self.d[key]
        context_meter.digest_metric("disk-read", 1, "count")
        context_meter.digest_metric("disk-read", len(pickled), "bytes")
        out = self.load(pickled)
        return out

    def __setitem__(self, key: Key, value: object) -> None:
        try:
            pickled = self.dump(value)
        except Exception as e:
            # zict.LRU ensures that the key remains in fast if we raise.
            # Wrap the exception so that it's recognizable by SpillBuffer,
            # which will then unwrap it.
            raise PickleError(key) from e

        # Thanks to Buffer.__setitem__, we never update existing
        # keys in slow, but always delete them and reinsert them.
        assert key not in self.d
        assert key not in self.weight_by_key

        pickled_size = sum(map(nbytes, pickled))

        if (
            self.max_weight is not False
            and self.total_weight.disk + pickled_size > self.max_weight
        ):
            # Stop callbacks and ensure that the key ends up in SpillBuffer.fast
            # To be caught by SpillBuffer.__setitem__
            raise MaxSpillExceeded(key)

        # Store to disk through File.
        # This may raise OSError, which is caught by SpillBuffer above.
        with context_meter.meter("disk-write", "seconds"):
            self.d[key] = pickled
        context_meter.digest_metric("disk-write", 1, "count")
        context_meter.digest_metric("disk-write", pickled_size, "bytes")

        weight = SpilledSize(safe_sizeof(value), pickled_size)
        self.weight_by_key[key] = weight
        self.total_weight += weight

    def __delitem__(self, key: Key) -> None:
        super().__delitem__(key)
        self.total_weight -= self.weight_by_key.pop(key)
