"""Per-request inference caches: (i) backbone per-layer KV, (ii) source-bank K/V projections. Isolated per
RequestBatch; workspace states are never stored across tokens (v1 contract)."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import torch

from .attention import LayerKVCache


@dataclass
class SourceBankCache:
    K: Optional[torch.Tensor] = None   # [B, t, w]
    V: Optional[torch.Tensor] = None

    @property
    def length(self) -> int:
        return 0 if self.K is None else self.K.shape[1]

    def append(self, K: torch.Tensor, V: torch.Tensor):
        if self.K is None:
            self.K, self.V = K, V
        else:
            self.K = torch.cat([self.K, K], 1)
            self.V = torch.cat([self.V, V], 1)

    def reset(self):
        self.K = self.V = None

    def nbytes(self) -> int:
        return 0 if self.K is None else (self.K.numel() + self.V.numel()) * self.K.element_size()


@dataclass
class RequestBatchCache:
    """One cache object per batch of concurrent requests (same length, no padding in v1)."""
    layers: list[LayerKVCache]
    bank: SourceBankCache = field(default_factory=SourceBankCache)

    @property
    def length(self) -> int:
        return self.layers[0].length

    def reset(self):
        for l in self.layers:
            l.reset()
        self.bank.reset()

    def nbytes(self) -> int:
        n = self.bank.nbytes()
        for l in self.layers:
            if l.k is not None:
                n += (l.k.numel() + l.v.numel()) * l.k.element_size()
        return n