"""The swappable attention contract and the registry (docs/ARCHITECTURE.md).

Every variant (mha, gqa, mla, gdn) subclasses `Attention`, registers itself under
a name, and is picked per layer through `ModelConfig.layer_types`. The rest of
the model never needs to know which one it got.
"""

from torch import Tensor, nn

from tinyllm.config import ModelConfig

ATTENTION: dict[str, type["Attention"]] = {}


def register(name: str):
    """Class decorator: @register("mha") adds the class to ATTENTION."""
    def deco(cls: type["Attention"]) -> type["Attention"]:
        if name in ATTENTION:
            raise ValueError(f"attention variant {name!r} already registered")
        ATTENTION[name] = cls
        return cls
    return deco


class Attention(nn.Module):
    def __init__(self, cfg: ModelConfig, layer_idx: int):
        super().__init__()
        self.cfg = cfg
        self.layer_idx = layer_idx

    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
        """x: (B, T, d_model) -> (B, T, d_model).

        rope: (cos, sin) for the T positions in x.
        cache: per-layer state for incremental decoding. None during training.
        """
        raise NotImplementedError

    def cache_bytes_per_token(self, dtype_bytes: int = 2) -> int:
        """Decode-time state this layer adds per generated token (bf16 = 2 bytes)."""
        raise NotImplementedError


def build_attention(cfg: ModelConfig, layer_idx: int) -> Attention:
    name = cfg.layer_types[layer_idx]
    if name not in ATTENTION:
        raise KeyError(f"unknown attention {name!r}; registered: {sorted(ATTENTION)}")
    return ATTENTION[name](cfg, layer_idx)