import torch
from typing import Optional

def compute_chunk_size(
    image: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
) -> int:
    """
    Determine a safe batch-size chunk for GPU processing based on the
    currently available CUDA memory.

    The estimation accounts for the memory required by the input image,
    an optional mask tensor, the output image, and temporary working
    tensors created during color conversion.

    Args:
        image (torch.Tensor):
            Input image tensor of shape (B, H, W, C).

        mask (torch.Tensor, optional):
            Optional mask tensor associated with the image batch.

    Returns:
        int:
            Maximum number of images that can be processed safely in a
            single chunk. If the image resides on the CPU, the complete
            batch size is returned.
    """
    if image.device.type != "cuda":
        return image.shape[0]

    free_mem, _ = torch.cuda.mem_get_info(image.device)

    image_batch = max(image.shape[0], 1)

    # Bytes required by one image sample
    image_sample_bytes = (
        image[0].numel() * image.element_size()
    )

    # Bytes required by one mask sample
    mask_sample_bytes = 0

    if mask is not None:
        if mask.device != image.device:
            mask = mask.to(image.device)

        mask_batch = max(mask.shape[0], 1)

        if mask_batch == image_batch:
            mask_sample_bytes = mask[0].numel() * mask.element_size()
        elif mask_batch == 1:
            # Single mask broadcast across the batch
            mask_sample_bytes = mask[0].numel() * mask.element_size()
        else:
            raise ValueError(
                f"Mask batch size ({mask_batch}) does not match image batch size ({image_batch})."
            )

    # Working memory:
    #   input image
    #   output image
    #   temporary color conversion tensors
    #   optional mask
    working_factor = 4

    required_per_sample = (
        image_sample_bytes * working_factor +
        mask_sample_bytes
    )

    # Keep a safety margin so other CUDA allocations have room.
    usable_memory = int(free_mem * 0.90)

    return max(1, min(image_batch, usable_memory // max(required_per_sample, 1)))