"""Cache adapters that manipulate real mlx_lm (0.32.x) KV tensors.

Dense mlx_lm KVCache stores post-RoPE keys/values [B, H_kv, T_buf, d_h] with buffer growth in steps of 256 and a
logical `offset`. QuantizedKVCache stores packed uint32 words + scales + biases per group (mx.quantize) and is consumed
by quantized_scaled_dot_product_attention (genuine quantized execution).

SparseKVCache / SparseQuantizedKVCache: keep a subset of positions (shared or per kv-head gather) while preserving
the *logical* offset so newly appended tokens get their original RoPE positions; retained keys already carry their
own positions (RoPE applied before caching). Multi-token appends use an explicit boolean mask: retained positions are
fully visible, new tokens are causal among themselves.
"""
from __future__ import annotations
from typing import Optional
import mlx.core as mx
from mlx.utils import tree_map
from mlx_lm.models.cache import KVCache, QuantizedKVCache


def dense_kv(c) -> tuple[mx.array, mx.array]:
    """Trimmed (offset-length) dense keys/values of a dense KVCache."""
    return c.keys_and_values()


def _bool_mask(n_new: int, n_kept: int) -> mx.array:
    left = mx.ones((n_new, n_kept), dtype=mx.bool_)
    r = mx.arange(n_new)
    right = r[:, None] >= r[None, :]
    return mx.concatenate([left, right], axis=1)


class SparseKVCache(KVCache):
    """KV cache holding a subset of positions with a preserved logical offset (dense dtype)."""

    def __init__(self):
        super().__init__()
        self.n_kept = 0
        self.retained_positions = None

    @classmethod
    def from_dense(cls, keys: mx.array, values: mx.array, logical_offset: int, indices=None) -> "SparseKVCache":
        c = cls()
        if indices is None:
            nk, nv = keys, values
            kept = int(keys.shape[2])
        else:
            idx = mx.array(indices, dtype=mx.int32)
            if idx.ndim == 1:
                nk = mx.take(keys, idx, axis=2); nv = mx.take(values, idx, axis=2); kept = int(idx.shape[0])
            elif idx.ndim == 2:
                g = idx[None, :, :, None]
                nk = mx.take_along_axis(keys, mx.broadcast_to(g, (keys.shape[0], keys.shape[1], idx.shape[1], keys.shape[3])), axis=2)
                nv = mx.take_along_axis(values, mx.broadcast_to(g, (values.shape[0], values.shape[1], idx.shape[1], values.shape[3])), axis=2)
                kept = int(idx.shape[1])
            else:
                raise ValueError("indices must be 1-D or 2-D")
            c.retained_positions = idx.tolist()
        c.keys, c.values = nk, nv
        c.n_kept = kept
        c.offset = int(logical_offset)
        mx.eval(c.keys, c.values)
        return c

    @classmethod
    def from_kvcache(cls, src: KVCache, indices=None, logical_offset: Optional[int] = None) -> "SparseKVCache":
        k, v = dense_kv(src)
        return cls.from_dense(k, v, src.offset if logical_offset is None else logical_offset, indices)

    def update_and_fetch(self, keys, values):
        n_new = keys.shape[2]
        if self.keys is None:
            self.keys, self.values = keys, values
        else:
            self.keys = mx.concatenate([self.keys[..., : self.n_kept, :], keys], axis=2)
            self.values = mx.concatenate([self.values[..., : self.n_kept, :], values], axis=2)
        self.n_kept += n_new
        self.offset += n_new
        return self.keys, self.values

    def keys_and_values(self):
        return self.keys[..., : self.n_kept, :], self.values[..., : self.n_kept, :]

    @property
    def state(self):
        return self.keys, self.values, self.offset, self.n_kept

    @state.setter
    def state(self, v):
        self.keys, self.values, self.offset, self.n_kept = v

    def size(self):
        return self.n_kept

    def make_mask(self, N, return_array=False, window_size=None):
        if N == 1:
            return None
        return _bool_mask(N, self.n_kept)

    def is_trimmable(self):
        return False

    @property
    def nbytes(self):
        if self.keys is None:
            return 0
        k, v = self.keys_and_values()
        return int(k.nbytes + v.nbytes)

    def to_quantized(self, group_size: int = 64, bits: int = 4) -> "SparseQuantizedKVCache":
        q = SparseQuantizedKVCache(group_size=group_size, bits=bits)
        q.offset = self.offset
        q.n_kept = self.n_kept
        k, v = self.keys_and_values()
        q.keys = mx.quantize(k, group_size=group_size, bits=bits)
        q.values = mx.quantize(v, group_size=group_size, bits=bits)
        mx.eval(q.keys, q.values)
        return q


class SparseQuantizedKVCache(QuantizedKVCache):
    """Packed quantized cache over retained positions with preserved logical offset."""

    def __init__(self, group_size: int = 64, bits: int = 8):
        super().__init__(group_size=group_size, bits=bits)
        self.n_kept = 0

    def update_and_fetch(self, keys, values):
        n_new = keys.shape[2]
        qk = mx.quantize(keys, group_size=self.group_size, bits=self.bits)
        qv = mx.quantize(values, group_size=self.group_size, bits=self.bits)
        if self.keys is None:
            self.keys, self.values = qk, qv
        else:
            cur_k, cur_v = self.keys_and_values()
            self.keys = tuple(mx.concatenate([a, b], axis=2) for a, b in zip(cur_k, qk))
            self.values = tuple(mx.concatenate([a, b], axis=2) for a, b in zip(cur_v, qv))
        self.n_kept += n_new
        self.offset += n_new
        return self.keys, self.values

    def keys_and_values(self):
        return (tree_map(lambda x: x[..., : self.n_kept, :], self.keys),
                tree_map(lambda x: x[..., : self.n_kept, :], self.values))

    @property
    def state(self):
        return self.keys, self.values, self.offset, self.group_size, self.bits, self.n_kept

    @state.setter
    def state(self, v):
        self.keys, self.values, self.offset, self.group_size, self.bits, self.n_kept = v

    def size(self):
        return self.n_kept

    def make_mask(self, N, return_array=False, window_size=None):
        if N == 1:
            return None
        return _bool_mask(N, self.n_kept)

    def is_trimmable(self):
        return False

    @property
    def nbytes(self):
        if self.keys is None:
            return 0
        k, v = self.keys_and_values()
        return int(sum(x.nbytes for x in k) + sum(x.nbytes for x in v))


# ---------------------------------------------------------------------------------------------------------------
def to_sparse_all(cache: list) -> list:
    """Wrap every dense layer into a SparseKVCache retaining all positions (100% retention control / B0 canonical form)."""
    return [SparseKVCache.from_kvcache(c) if type(c) is KVCache else c for c in cache]


def quantize_cache(cache: list, bits: int, group_size: int = 64) -> list:
    """B2: packed quantization of every layer's K/V (storage AND execution representation)."""
    out = []
    for c in cache:
        if isinstance(c, SparseKVCache):
            out.append(c.to_quantized(group_size=group_size, bits=bits))
        elif type(c) is KVCache:
            out.append(SparseKVCache.from_kvcache(c).to_quantized(group_size=group_size, bits=bits))
        else:
            raise TypeError(f"cannot quantize cache of type {type(c)}")
    return out


def select_positions(cache: list, indices_per_layer, logical_offset: Optional[int] = None) -> list:
    """B3/B4/B5/B6: sparse cache retaining the given positions per layer (1-D shared or 2-D per kv-head indices)."""
    out = []
    for c, idx in zip(cache, indices_per_layer):
        if isinstance(c, SparseKVCache):
            k, v = c.keys_and_values()
            out.append(SparseKVCache.from_dense(k, v, c.offset if logical_offset is None else logical_offset, idx))
        elif type(c) is KVCache:
            out.append(SparseKVCache.from_kvcache(c, idx, logical_offset))
        else:
            raise TypeError("select_positions expects dense layers")
    return out


def copy_cache(cache: list) -> list:
    """Independent branch of a cache (no shared mutable arrays)."""
    out = []
    for c in cache:
        if isinstance(c, SparseQuantizedKVCache):
            n = SparseQuantizedKVCache(group_size=c.group_size, bits=c.bits)
            k, v = c.keys_and_values()
            n.keys = tuple(mx.array(x) for x in k); n.values = tuple(mx.array(x) for x in v)
            n.offset, n.n_kept = c.offset, c.n_kept
        elif isinstance(c, SparseKVCache):
            k, v = c.keys_and_values()
            n = SparseKVCache(); n.keys, n.values = mx.array(k), mx.array(v); n.offset, n.n_kept = c.offset, c.n_kept
            n.retained_positions = c.retained_positions
        elif type(c) is KVCache:
            k, v = c.keys_and_values()
            n = KVCache(); n.keys, n.values = mx.array(k), mx.array(v); n.offset = c.offset
        elif type(c) is QuantizedKVCache:
            k, v = c.keys_and_values()
            n = QuantizedKVCache(group_size=c.group_size, bits=c.bits)
            n.keys = tuple(mx.array(x) for x in k); n.values = tuple(mx.array(x) for x in v); n.offset = c.offset
        else:
            raise TypeError(type(c))
        out.append(n)
    mx.eval([x for c in out for x in _arrays_of(c)])
    return out


def _arrays_of(c) -> list[mx.array]:
    k, v = c.keys_and_values()
    if isinstance(k, tuple):
        return list(k) + list(v)
    return [k, v]


def cache_tensor_bytes(cache: list) -> int:
    """Stored tensor bytes over the *logical content* (trimmed to stored positions, not allocator buffers)."""
    return int(sum(c.nbytes if not (type(c) is KVCache or type(c) is QuantizedKVCache) else sum(x.nbytes for x in _arrays_of(c)) for c in cache))


def dense_bytes_formula(n_layers: int, T: int, n_kv_heads: int, head_dim: int, bytes_per_scalar: int, B: int = 1) -> int:
    """M_KV = 2 B L T H_kv d_h s (consistency check only; measured bytes are authoritative)."""
    return 2 * B * n_layers * T * n_kv_heads * head_dim * bytes_per_scalar