"""
This type stub file was generated by pyright.
"""

from typing import Any, Dict, List, Optional, Protocol, Literal, Self

import mlx.nn as nn
from mlx.core import array
import mlx.core as mx

class Cache(Protocol):
    keys: mx.array
    values: mx.array
    offset: int
    def update_and_fetch(
        self, keys: mx.array, values: mx.array
    ) -> tuple[mx.array, mx.array]: ...
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v) -> None: ...

def make_prompt_cache(
    model: nn.Module, max_kv_size: Optional[int] = ...
) -> List[Cache | Any]:
    """
    Construct the model's cache for use in generation.

    This function will defer the cache construction to the model if it has a
    ``make_cache`` method, otherwise it will make a default KV cache.

    Args:
        model (nn.Module): The language model.
        max_kv_size (Optional[int]): If provided and the model does not have a
            ``make_cache`` method, a ``RotatingKVCache`` is used with a maximum
            size of ``max_kv_size``
    """

def save_prompt_cache(
    file_name: str, cache: List[Cache], metadata: Dict[str, str] = ...
) -> None:
    """
    Save a pre-computed prompt cache to a file.

    Args:
        file_name (str): The ``.safetensors`` file name.
        cache (List[Any]): The model state.
        metadata (Dict[str, str]): Optional metadata to save along with model
            state.
    """

def load_prompt_cache(file_name: str, return_metadata=...) -> array:
    """
    Load a prompt cache from a file.

    Args:
        file_name (str): The ``.safetensors`` file name.
        return_metadata (bool): Whether or not to return metadata.
            Default: ``False``.

    Returns:
        List[Any] or Tuple[List[Any], Dict[str, str]]: The prompt cache and
            the metadata if requested.
    """

def can_trim_prompt_cache(cache: List[Cache]) -> bool:
    """
    Check if model's cache can be trimmed.
    """

def trim_prompt_cache(cache: List[Cache], num_tokens: int) -> List[Cache]:
    """
    Trim the model's cache by the given number of tokens.

    This function will trim the cache if possible (in-place) and return the
    number of tokens that were trimmed.

    Args:
        cache (List[Any]): The model's cache.
        num_tokens (int): The number of tokens to trim.

    Returns:
        (int): The number of tokens that were trimmed.
    """

def create_attention_mask(
    N: int, offset: int, return_array: bool, window_size: Optional[int]
) -> array | Literal["causal"] | None: ...

class _BaseCache(Cache):
    keys: mx.array | None
    values: mx.array | None
    offset: int
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v) -> None: ...
    @property
    def meta_state(self) -> Literal[""]: ...
    @meta_state.setter
    def meta_state(self, v) -> None: ...
    def trim(self, n: int) -> int: ...
    def is_trimmable(self) -> Literal[False]: ...
    @classmethod
    def from_state(cls, state, meta_state) -> Self: ...

class ConcatenateKVCache(_BaseCache):
    """ConcatenateKVCache the simplest KV cache implementation.

    Can be used as a mock KV cache or when large blocks are being processed at
    a time in which case KVCache isn't necessarily faster. Consider using the
    KVCache with a larger step size before using this cache.
    """
    def __init__(self) -> None: ...
    def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
        ...
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    def is_trimmable(self):  # -> Literal[True]:
        ...
    def trim(self, n: int) -> int: ...
    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
        ...

class QuantizedKVCache(_BaseCache):
    step = ...
    def __init__(self, group_size: int = ..., bits: int = ...) -> None: ...
    def update_and_fetch(self, keys, values):  # -> Any:
        ...
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    @property
    def meta_state(self):  # -> tuple[str, ...]:
        ...
    @meta_state.setter
    def meta_state(self, v):  # -> None:
        ...
    def is_trimmable(self):  # -> Literal[True]:
        ...
    def trim(self, n: int) -> int: ...
    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
        ...

class KVCache(_BaseCache):
    step = ...
    def __init__(self) -> None: ...
    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
        ...
    @property
    def state(
        self,
    ) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v) -> None: ...
    def is_trimmable(self):  # -> Literal[True]:
        ...
    def trim(self, n: int) -> int: ...
    def to_quantized(
        self, group_size: int = ..., bits: int = ...
    ) -> QuantizedKVCache: ...
    def make_mask(
        self, *args: Any, **kwargs: Any
    ) -> mx.array | Literal["causal"] | None: ...

class RotatingKVCache(_BaseCache):
    step = ...
    keys: mx.array | None
    values: mx.array | None
    keep: int
    max_size: int
    _idx: int
    def __init__(self, max_size, keep=...) -> None: ...
    def _trim(
        self, trim_size: int, v: mx.array, append: mx.array | None = ...
    ) -> mx.array: ...
    def update_and_fetch(
        self, keys, values
    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
        ...
    @property
    def state(
        self,
    ) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    @property
    def meta_state(self):  # -> tuple[str, ...]:
        ...
    @meta_state.setter
    def meta_state(self, v):  # -> None:
        ...
    def is_trimmable(self):  # -> bool:
        ...
    def trim(self, n: int) -> int: ...
    def to_quantized(
        self, group_size: int = ..., bits: int = ...
    ) -> QuantizedKVCache: ...
    def make_mask(
        self, N: int, window_size: Optional[int] = ..., return_array: bool = ...
    ):  # -> array | Literal['causal'] | None:
        ...

class ArraysCache(_BaseCache):
    def __init__(self, size, left_padding: Optional[List[int]] = ...) -> None: ...
    def __setitem__(self, idx, value):  # -> None:
        ...
    def __getitem__(self, idx): ...
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    def filter(self, batch_indices):  # -> None:
        """
        In-place filter to keep just the given indices in the cache.
        """

    def extend(self, other):  # -> None:
        """
        In-place extend this cache with the other cache.
        """

    def make_mask(self, N: int) -> mx.array | None: ...

class MambaCache(ArraysCache):
    def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...

class ChunkedKVCache(KVCache):
    def __init__(self, chunk_size) -> None: ...
    def maybe_trim_front(self):  # -> None:
        ...
    def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
        ...
    def trim(self, n: int) -> int: ...
    @property
    def meta_state(self):  # -> tuple[str, ...]:
        ...
    @meta_state.setter
    def meta_state(self, v):  # -> None:
        ...

class CacheList(_BaseCache):
    def __init__(self, *caches) -> None: ...
    def __getitem__(self, idx): ...
    def is_trimmable(self):  # -> bool:
        ...
    def trim(self, n: int) -> int: ...
    @property
    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    def filter(self, batch_indices):  # -> None:
        """
        In-place filter to keep just the given indices in the cache.
        """

    def extend(self, other):  # -> None:
        """
        In-place extend this cache with the other cache.
        """

class BatchKVCache(_BaseCache):
    step: int
    keys: array | None
    values: array | None
    offset: array
    left_padding: array
    _idx: int
    def __init__(self, left_padding: List[int]) -> None: ...
    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
    @property
    def state(
        self,
    ):  # -> tuple[Any | array | None, Any | array | None, array | Any, array | Any]:
        ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    def is_trimmable(self):  # -> Literal[True]:
        ...
    def trim(self, n):  # -> int | float:
        ...
    def make_mask(self, N: int, return_array: bool = ..., **kwargs):  # -> array:
        ...
    def filter(self, batch_indices):  # -> None:
        """
        In-place filter to keep just the given indices in the cache.
        """

    def extend(self, other):  # -> None:
        """
        In-place extend this cache with the other cache.
        """

class BatchRotatingKVCache(_BaseCache):
    step: int
    keys: array | None
    values: array | None
    offset: array
    left_padding: array
    max_size: int
    _idx: int
    _offset: int
    rotated: bool
    _lengths: array | None
    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
    @property
    def state(
        self,
    ):  # -> tuple[Any | array | None, Any | array | None, array | Any, array | Any]:
        ...
    @state.setter
    def state(self, v):  # -> None:
        ...
    @property
    def meta_state(self):  # -> tuple[str, ...]:
        ...
    @meta_state.setter
    def meta_state(self, v):  # -> None:
        ...
    def is_trimmable(self):  # -> bool:
        ...
    def trim(self, n):  # -> int:
        ...
    def to_quantized(
        self, group_size: int = ..., bits: int = ...
    ) -> QuantizedKVCache: ...
    def make_mask(
        self, N: int, window_size: Optional[int] = ..., return_array: bool = ...
    ):  # -> array:
        ...
    def filter(self, batch_indices):  # -> None:
        """
        In-place filter to keep just the given indices in the cache.
        """

    def extend(self, other):  # -> None:
        """
        In-place extend this cache with the other cache.
        """
