"""Build the text and visual prefix used for waypoint prediction."""

from __future__ import annotations

from typing import List, Optional, Tuple

import torch
import torch.nn as nn

from .history_tokens import interleave_tvi_tokens


def build_visual_prefix(
    llm: nn.Module,
    tokenizer,
    projector: nn.Module,
    tvi: nn.Module,
    coarse_tokens: torch.Tensor,
    coarse_tidx: torch.Tensor,
    fine_tokens: torch.Tensor,
    fine_tidx: torch.Tensor,
    instructions: List[str],
    history_valid_mask: torch.Tensor,
    *,
    yaw_hist: Optional[torch.Tensor] = None,
    yaw_curr: Optional[torch.Tensor] = None,
    use_angle: bool = False,
    include_history: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Build the text/visual prefix before the ACT token.

    The current observation is always represented by ``fine_tokens``.  Callers
    must provide past observations separately as ``coarse_tokens``; an empty
    history is represented by zero-filled coarse slots and a false validity
    mask, rather than by duplicating the current image.
    """
    if coarse_tokens.ndim != 3 or fine_tokens.ndim != 3:
        raise ValueError("coarse_tokens and fine_tokens must be [B, N, C]")
    batch_size = fine_tokens.size(0)
    if coarse_tokens.size(0) != batch_size:
        raise ValueError("coarse and fine batch sizes must match")
    if (history_valid_mask is None or history_valid_mask.dtype != torch.bool
            or history_valid_mask.ndim != 2):
        raise ValueError("history_valid_mask must be a boolean [B, H] tensor")
    if history_valid_mask.size(0) != batch_size:
        raise ValueError("history_valid_mask batch size must match visual tokens")

    device = next(projector.parameters()).device
    vis_fine = projector(fine_tokens.to(device))
    pieces = []
    masks = []
    if include_history:
        vis_coarse = projector(coarse_tokens.to(device))
        vis_coarse, coarse_attention = interleave_tvi_tokens(
            tvi,
            vis_coarse,
            coarse_tidx.to(device),
            kind_id=0,
            yaw_per_frame=yaw_hist,
            use_angle=use_angle,
            frame_valid_mask=history_valid_mask.to(device),
        )
        pieces.append(vis_coarse)
        masks.append(coarse_attention.to(dtype=torch.long))
    vis_fine, fine_attention = interleave_tvi_tokens(
        tvi,
        vis_fine,
        fine_tidx.to(device),
        kind_id=1,
        yaw_per_frame=yaw_curr,
        use_angle=use_angle,
    )
    tok = tokenizer(
        instructions,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128,
    )
    tok = {key: value.to(device) for key, value in tok.items()}
    text_emb = llm.get_input_embeddings()(tok["input_ids"])
    pieces.append(vis_fine)
    masks.extend([fine_attention.to(dtype=torch.long)])
    sequence = torch.cat([text_emb, *pieces], dim=1).to(llm.dtype)
    attention = torch.cat([tok["attention_mask"], *masks], dim=1)
    return sequence, attention