"""Ephemeral, branchable inference state; never stored in training checkpoints.

Treat returned tensors as read-only. Appending to a cache forks its containers
and replaces tensors, so an answer branch cannot contaminate another answer.
Weight sharing does NOT imply cache sharing: every depth/pass has its own slot.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
import weakref

import torch


@dataclass
class AttentionState:
    key: torch.Tensor | None = None
    value: torch.Tensor | None = None


@dataclass
class WorkspaceState:
    owner: weakref.ReferenceType
    signature: tuple
    passes: int
    batch_size: int
    length: int = 0
    tonic: torch.Tensor | None = None
    previous_embedding: torch.Tensor | None = None
    previous_token: torch.Tensor | None = None
    attention: dict[str, AttentionState] = field(default_factory=dict)

    def fork(self):
        """Share immutable tensor storage; copy all subsequently changed holders."""
        return replace(self, attention={key: replace(value) for key, value in self.attention.items()})

    def select(self, rows):
        """Select/repeat equal-length prompt rows for independent continuations."""
        if self.previous_token is None:
            raise ValueError("cannot select an empty cache")
        indices = torch.as_tensor(rows, dtype=torch.long, device=self.previous_token.device)
        if indices.ndim != 1 or not indices.numel():
            raise ValueError("cache selection requires nonempty one-dimensional rows")
        result = self.fork()
        result.batch_size = indices.numel()
        result.tonic = self.tonic.index_select(1, indices)
        result.previous_embedding = self.previous_embedding.index_select(0, indices)
        result.previous_token = self.previous_token.index_select(0, indices)
        result.attention = {key: AttentionState(value.key.index_select(0, indices),
                                               value.value.index_select(0, indices))
                            for key, value in self.attention.items()}
        return result


def inference_signature(model, device):
    # Detect optimizer/load_state_dict/in-place edits and dtype/device changes.
    # The EMA embedding is training-only and does not affect cached inference.
    tensors = list(model.parameters()) + [b for name, b in model.named_buffers() if name != "ema_embedding"]
    versions = tuple((id(t), t._version, t.device, t.dtype) for t in tensors)
    autocast = torch.is_autocast_enabled(device.type)
    return (versions, repr(model.cfg), device, autocast,
            torch.get_autocast_dtype(device.type) if autocast else None)


def prepare_state(model, ids, passes, cache):
    if model.training or torch.is_grad_enabled():
        raise ValueError("cached inference requires model.eval() and no_grad/inference_mode")
    signature = inference_signature(model, ids.device)
    if cache is None:
        return WorkspaceState(weakref.ref(model), signature, passes, ids.shape[0])
    if not isinstance(cache, WorkspaceState) or cache.owner() is not model:
        raise ValueError("cache belongs to another model")
    if cache.signature != signature:
        raise ValueError("cache is stale: weights, configuration, device, or autocast changed")
    if cache.passes != passes or cache.batch_size != ids.shape[0]:
        raise ValueError("cache pass count or batch size mismatch; use cache.select for branching")
    if cache.length + ids.shape[1] > model.cfg.context:
        raise ValueError("cache exceeds context; explicitly recompute a bounded context window")
    return cache.fork()