from __future__ import annotations

import torch

from .recurrent.positions import build_compact_positions


@torch.no_grad()
def greedy_generate(
    model,
    batch: dict[str, torch.Tensor],
    max_new_tokens: int = 64,
    eos_token_id: int | None = None,
) -> torch.LongTensor:
    """Correctness-first generation that recomputes the recurrent sequence each step."""
    if batch["input_ids"].shape[0] != 1:
        raise ValueError("greedy_generate currently supports batch_size=1")
    model.eval()
    input_ids = batch["input_ids"]
    header_length = int(batch["header_length"])
    instruction_length = int(batch["instruction_length"])
    visual_token_count = int(batch["visual_token_counts"][0])
    grid_height = int(batch["grid_heights"][0])
    grid_width = int(batch["grid_widths"][0])
    visual_embeddings = model.encode_images(
        batch["pixel_values"], batch["grid_sizes"], batch["merge_sizes"]
    )
    generated: list[torch.Tensor] = []
    for _ in range(max_new_tokens):
        response_length = input_ids.shape[1] - header_length - visual_token_count - instruction_length
        mask_positions, rope_positions, visual_mask, instruction_mask = build_compact_positions(
            torch.tensor([header_length], device=input_ids.device),
            torch.tensor([instruction_length], device=input_ids.device),
            torch.tensor([response_length], device=input_ids.device),
            visual_token_counts=torch.tensor([visual_token_count], device=input_ids.device),
            grid_heights=torch.tensor([grid_height], device=input_ids.device),
            grid_widths=torch.tensor([grid_width], device=input_ids.device),
            grid_h=model.config.vision.grid_height,
            grid_w=model.config.vision.grid_width,
            rope_mode=model.config.recurrent_visual.rope_mode,
        )
        text_embeddings = model.language_model.get_input_embeddings()(input_ids)
        inputs_embeds = model._replace_visual_tokens(text_embeddings, visual_embeddings, visual_mask)
        prefix_length = header_length + visual_token_count + instruction_length
        token_type_ids = torch.zeros_like(input_ids)
        token_type_ids[:, :prefix_length] = 1
        outputs = model.recurrent_model(
            inputs_embeds=inputs_embeds,
            attention_mask=torch.ones_like(input_ids),
            token_type_ids=token_type_ids,
            position_ids=mask_positions,
            rope_position_ids=rope_positions,
            visual_mask=visual_mask,
            instruction_mask=instruction_mask,
            use_cache=False,
        )
        next_token = model.language_model.lm_head(outputs.last_hidden_state[:, -1]).argmax(dim=-1)
        generated.append(next_token)
        input_ids = torch.cat((input_ids, next_token[:, None]), dim=1)
        if eos_token_id is not None and bool((next_token == eos_token_id).all()):
            break
    return torch.stack(generated, dim=1) if generated else input_ids[:, :0]