"""Autoregressive sampling (no KV cache)."""

from __future__ import annotations

from collections.abc import Callable

import torch

from zenith.model.model import ZenithLM


@torch.no_grad()
def generate(
    model: ZenithLM,
    idx: torch.Tensor,
    max_new_tokens: int,
    temperature: float = 1.0,
    top_k: int | None = None,
    eos_id: int | None = None,
    generator: torch.Generator | None = None,
    logits_hook: Callable[[list[int], torch.Tensor], torch.Tensor] | None = None,
    stop_fn: Callable[[list[int]], bool] | None = None,
) -> torch.Tensor:
    """Extend `idx` (shape (1, T)) by up to `max_new_tokens`; temperature 0 is greedy.

    `logits_hook(ids, logits)` may rewrite the next-token logits (e.g. to mask tokens);
    `stop_fn(ids)` ends generation after the token for which it returns True.
    """
    for _ in range(max_new_tokens):
        logits, _ = model(idx[:, -model.cfg.context_len :])
        logits = logits[:, -1, :].float()
        if logits_hook is not None:
            logits = logits_hook(idx[0].tolist(), logits)
        if temperature == 0:
            next_id = logits.argmax(dim=-1, keepdim=True)
        else:
            logits = logits / temperature
            if top_k is not None:
                kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
                logits = logits.masked_fill(logits < kth, float("-inf"))
            probs = logits.softmax(dim=-1)
            next_id = torch.multinomial(probs.cpu(), 1, generator=generator).to(idx.device)
        idx = torch.cat([idx, next_id], dim=1)
        if eos_id is not None and next_id.item() == eos_id:
            break
        if stop_fn is not None and stop_fn(idx[0].tolist()):
            break
    return idx