import torch
import random

MASK_TOKEN = 0
PAD_TOKEN = 1
CLS_TOKEN = 2


def complete_masking(batch, masking_p, n_tokens):
    """Apply masking to input batch for masked language modeling.

    Args:
        batch (dict): Input batch containing 'input_ids' and 'attention_mask'
        masking_p (float): Probability of masking a token
        n_tokens (int): Total number of tokens in vocabulary

    Returns:
        dict: Batch with masked indices and masking information
    """
    device = batch["input_ids"].device
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]

    # Create mask tensor (1 for tokens to be masked, 0 otherwise)
    prob = torch.rand(input_ids.shape, device=device)
    mask = (prob < masking_p) & (input_ids != PAD_TOKEN) & (input_ids != CLS_TOKEN)

    # For masked tokens:
    # - 80% replace with MASK token
    # - 10% replace with random token
    # - 10% keep unchanged
    masked_indices = input_ids.clone()

    # Calculate number of tokens to be masked
    num_tokens_to_mask = mask.sum().item()

    # Determine which tokens get which type of masking
    mask_mask = torch.rand(num_tokens_to_mask, device=device) < 0.8
    random_mask = (torch.rand(num_tokens_to_mask, device=device) < 0.5) & ~mask_mask

    # Apply MASK token (80% of masked tokens)
    masked_indices[mask] = torch.where(
        mask_mask,
        torch.tensor(MASK_TOKEN, device=device, dtype=torch.long),
        masked_indices[mask],
    )

    # Apply random tokens (10% of masked tokens).
    # Build a full-shape boolean mask first; single-level boolean assignment
    # modifies the tensor in-place correctly (double-index does not).
    random_tokens = torch.randint(
        3,
        n_tokens,  # Start from 3 to avoid special tokens
        (random_mask.sum(),),
        device=device,
        dtype=torch.long,
    )
    random_full_mask = torch.zeros_like(input_ids, dtype=torch.bool)
    random_full_mask[mask] = random_mask
    masked_indices[random_full_mask] = random_tokens

    # 10% remain unchanged

    return {
        "masked_indices": masked_indices,
        "attention_mask": attention_mask,
        "mask": mask,
        "input_ids": input_ids,
    }
