from pathlib import Path
from typing import Any, Union

from vidata.writers.base_writer import BaseWriter

PathLike = Union[str, Path]


class MultilabelWriter(BaseWriter):
    """Writer for multilabel mask data."""

    role: str = "mask"


class MultilabelStackedWriter(BaseWriter):
    """Writer for stacked multilabel mask data.

    Each channel of the input tensor is saved as an individual file.
    The filenames are generated by appending a zero-padded class index to the base filename.

    Attributes:
        role (str): Type of the data role, fixed as `"mask"`.
        num_classes (int): Number of classes (channels) in the mask data.
        zero_padding (int): Number of digits to use for zero-padding the class index in filenames.
    """

    role: str = "mask"

    def __init__(self, ftype: str, num_classes: int, backend=None, zero_padding=4):
        super().__init__(ftype, backend)
        self.num_classes = num_classes
        self.zero_padding = zero_padding

    def save(self, data: Any, file: PathLike, *args, **kwargs) -> None:
        """Save each channel of a multilabel mask as a separate file.

        Args:
            data (Union[np.ndarray, Sequence[np.ndarray]]): Mask data with shape `(num_classes, H, W, ...)`.
                The first dimension must match `num_classes`.
            file (str): Base filename without extension or class index.
            *args: Additional positional arguments passed to `save_file`.
            **kwargs: Additional keyword arguments passed to `save_file`.

        Raises:
            AssertionError: If the number of channels in `data` does not match `num_classes`.
        """
        files = [
            f"{file}_{str(i).zfill(self.zero_padding)}{self.ftype}" for i in range(self.num_classes)
        ]
        assert (
            len(files) == data.shape[0]
        ), "Number of channels must match the first dimension of data."
        for i, file_i in enumerate(files):
            self.save_file(data[i, ...], file_i, *args, **kwargs)
