"""Batched greedy generation for benchmarks (no KV cache: simple, not fast)."""
import torch

from ar.model.mask import build_doc_block_mask


@torch.no_grad()
def generate_batch(model, tok, prompts, seq_len, max_new, stop, device):
    """Greedy-decode a batch of prompts; each output is cut at its first stop string."""
    B = len(prompts)
    # fixed-size buffer: each row has its own prompt length and write position;
    # the causal mask hides the placeholder slots to the right of each position
    buf = torch.full((B, seq_len), tok.eot_id, dtype=torch.long, device=device)
    pos = []
    for b, p in enumerate(prompts):
        x = tok.encode(p)[-(seq_len - max_new):]   # keep the end if a prompt is too long
        buf[b, :len(x)] = torch.tensor(x, device=device)
        pos.append(len(x))
    start = list(pos)
    seg = torch.zeros(B, seq_len, dtype=torch.int32, device=device)   # one doc per row
    mask = build_doc_block_mask(seg, seq_len, device=device)
    done, outs = [False] * B, [""] * B
    rows = torch.arange(B, device=device)

    for _ in range(max_new):
        logits = model(buf, mask)
        # each row's next token = most likely token at its own current position
        idx = torch.tensor(pos, device=device) - 1
        nxt = logits[rows, idx].argmax(-1).tolist()
        for b in range(B):
            if done[b]:
                continue
            if nxt[b] == tok.eot_id:                 # model ended the text
                done[b] = True
                continue
            buf[b, pos[b]] = nxt[b]
            pos[b] += 1
            outs[b] = tok.decode(buf[b, start[b]:pos[b]].tolist())
            if pos[b] >= seq_len or any(s in outs[b] for s in stop):
                done[b] = True
        if all(done):
            break

    # cut each output at its first stop string
    for b in range(B):
        hits = [outs[b].find(s) for s in stop if s in outs[b]]
        outs[b] = outs[b][: min(hits)] if hits else outs[b]
    return outs