"""Inference-only generation, independent per-request caches and EOS handling."""
import torch
import torch.nn.functional as F


def validate_inputs(ids, attention_mask, context):
    if ids is None or ids.ndim != 2 or not ids.shape[0] or not ids.shape[1]:
        raise ValueError("Expected a nonempty [batch, sequence] input_ids tensor")
    if ids.dtype not in (torch.int32, torch.int64):
        raise ValueError("input_ids must contain integer token IDs")
    if attention_mask is not None and (attention_mask.shape != ids.shape or attention_mask.device != ids.device):
        raise ValueError("attention_mask shape and device must match input_ids")
    if attention_mask is not None and not torch.all((attention_mask == 0) | (attention_mask == 1)):
        raise ValueError("attention_mask must contain only zero and one")
    mask = torch.ones_like(ids, dtype=torch.bool) if attention_mask is None else attention_mask.bool()
    for row in mask:
        pos = row.nonzero(as_tuple=True)[0]
        if not len(pos) or int(pos[-1]-pos[0]+1) != len(pos):
            raise ValueError("Only contiguous nonempty prompts with left/right padding are supported")
        if len(pos) > context:
            raise ValueError("Prompt exceeds the model context")
    return mask


def sample_token(logits, do_sample, temperature=1.0, top_p=1.0, top_k=0):
    if not do_sample:
        return logits.argmax(dim=-1)
    if temperature <= 0 or not 0 < top_p <= 1 or type(top_k) is not int or top_k < 0:
        raise ValueError("Invalid sampling parameters")
    scores = logits.float() / temperature
    if top_k:
        cutoff = scores.topk(min(top_k, scores.shape[-1]), dim=-1).values[:, -1:]
        scores = scores.masked_fill(scores < cutoff, -torch.inf)
    if top_p < 1:
        ordered, indices = scores.sort(descending=True, dim=-1)
        remove = ordered.softmax(-1).cumsum(-1) > top_p
        remove[:, 1:] = remove[:, :-1].clone()
        remove[:, 0] = False
        scores = scores.masked_fill(torch.zeros_like(remove).scatter(1, indices, remove), -torch.inf)
    return torch.multinomial(scores.softmax(-1), 1).squeeze(-1)


def prompt_batches(ids, mask, limit):
    lengths = mask.sum(-1).tolist()
    start = 0
    while start < len(lengths):
        end = start + 1
        while end < len(lengths) and end-start < limit and lengths[end] == lengths[start]:
            end += 1
        yield torch.stack([ids[i, mask[i]] for i in range(start, end)])
        start = end


@torch.no_grad()
def generate_batched(model, ids, mask, requested, eos_ids, pad, option, limit):
    decoder = None
    if model.config.fast_decode:
        if model.config.inference_backend in ("torch", "cpu"):
            decoder = (lambda backbone, **kw: backbone.enable_decode(**kw),
                       lambda backbone: backbone.disable_decode())
        else:
            from .fast_decode import enable, disable
            decoder = (enable, disable)
    pad = pad if pad is not None else min(eos_ids, default=0)
    eos = torch.tensor(sorted(eos_ids), device=ids.device, dtype=ids.dtype)
    was_training = model.training
    completions = []
    model.eval()
    try:
        for prompt in prompt_batches(ids, mask, limit):
            remaining = model.config.max_position_embeddings-prompt.shape[1]
            if remaining < 1 or (requested > remaining and model.config.generation_context_policy == "error"):
                raise ValueError("Requested generation exceeds the model context")
            budget = min(requested, remaining)
            count = prompt.shape[0]
            tail = ids.new_full((count, budget), pad)
            alive = torch.ones(count, device=ids.device, dtype=torch.bool)
            lengths = torch.zeros(count, device=ids.device, dtype=torch.long)
            try:
                if decoder:
                    decoder[0](model.backbone, max_seq_len=prompt.shape[1]+budget, batch_size=count)
                current = prompt
                for step in range(budget):
                    token = sample_token(model._backbone_logits(current)[:, -1], option("do_sample", False),
                        option("temperature", 1.0), option("top_p", 1.0), option("top_k", 0))
                    token = torch.where(alive, token, pad)
                    tail[:, step] = token
                    lengths += alive
                    if eos.numel():
                        alive &= ~(token[:, None] == eos[None, :]).any(-1)
                    if not alive.any():
                        break
                    current = token[:, None] if decoder else torch.cat([current, token[:, None]], dim=1)
            finally:
                if decoder:
                    decoder[1](model.backbone)
            completions.extend(tail[row, :n] for row, n in enumerate(lengths.tolist()))
    finally:
        model.train(was_training)
    width = max(map(len, completions))
    tails = torch.stack([F.pad(x, (0, width-len(x)), value=pad) for x in completions])
    return torch.cat([ids, tails], dim=1)