"""Schema for exact all-layer BF16 Teacher Text K/V caches."""

from __future__ import annotations

from typing import Mapping

import torch


SCHEMA_VERSION_TEXT_KV_ALL54 = "hyworldplay_teacher_text_kv_exact_all54_bf16_v1"
TEXT_KV_LAYER_COUNT = 54
TEXT_KV_HEADS = 16
TEXT_KV_TOKENS = 903
TEXT_KV_HEAD_DIM = 128


def text_kv_key(block_id: int, kind: str) -> str:
    if not 0 <= block_id < TEXT_KV_LAYER_COUNT:
        raise ValueError(f"Invalid Text K/V block: {block_id}")
    if kind not in {"k_txt", "v_txt"}:
        raise ValueError(f"Invalid Text K/V kind: {kind}")
    return f"block_{block_id:02d}_{kind}"


def expected_text_kv_keys() -> set[str]:
    return {
        text_kv_key(block_id, kind)
        for block_id in range(TEXT_KV_LAYER_COUNT)
        for kind in ("k_txt", "v_txt")
    } | {"text_valid_mask"}


def validate_text_kv_all54(tensors: Mapping[str, torch.Tensor]) -> None:
    missing = expected_text_kv_keys().difference(tensors)
    unexpected = set(tensors).difference(expected_text_kv_keys())
    if missing or unexpected:
        raise ValueError(
            f"Text K/V key mismatch: missing={sorted(missing)}, "
            f"unexpected={sorted(unexpected)}"
        )
    expected_shape = (1, TEXT_KV_HEADS, TEXT_KV_TOKENS, TEXT_KV_HEAD_DIM)
    for block_id in range(TEXT_KV_LAYER_COUNT):
        for kind in ("k_txt", "v_txt"):
            name = text_kv_key(block_id, kind)
            value = tensors[name]
            if tuple(value.shape) != expected_shape:
                raise ValueError(f"Unexpected {name} shape: {tuple(value.shape)}")
            if value.dtype != torch.bfloat16:
                raise TypeError(f"{name} must be BF16, got {value.dtype}")
            if not torch.isfinite(value).all():
                raise ValueError(f"{name} contains non-finite values")
    mask = tensors["text_valid_mask"]
    if tuple(mask.shape) != (1, TEXT_KV_TOKENS) or mask.dtype != torch.bool:
        raise ValueError(
            f"Unexpected text_valid_mask: shape={tuple(mask.shape)}, dtype={mask.dtype}"
        )