import json
from copy import deepcopy
from importlib.resources import files

from .. import RELAIData
from .image import Image as RELAIImage


class ImageGrid(RELAIData):
    """Data model for an image grid component.

    Args:
        images (list[Image]): List of RELAI images to be displayed in the grid.
        grid_size (tuple[int, int]): Size of the grid in terms of rows and columns.
        row_captions (list[str], optional): List of row captions. Defaults to None.
        column_captions (list[str], optional): List of column captions. Defaults to None.
        caption (str, optional): Caption for the image grid. Defaults to None.
        columnSpan (int, optional): Number of columns the component should span. Defaults to 12.
    """

    schema_file: str = files("relai_visualization").joinpath("schema/components/image-grid.schema.json")

    def __init__(
        self,
        images: list[RELAIImage],
        grid_size: tuple[int, int],
        row_captions: list[str] = None,
        column_captions: list[str] = None,
        caption: str = None,
        column_span: int = 12,
    ) -> None:
        if len(grid_size) != 2 or not all(isinstance(i, int) for i in grid_size):
            raise ValueError("Grid size must be a tuple of two integers")

        if grid_size[0] < 1 or grid_size[1] < 1:
            raise ValueError("Grid parameters must be greater than 0")

        if row_captions and len(row_captions) < 1:
            raise ValueError("Row captions cannot be empty when provided.")

        if column_captions and len(column_captions) < 1:
            raise ValueError("Column captions cannot be empty when provided.")

        self._image_objects = images
        images = [image.data for image in self._image_objects]
        grid_size = list(grid_size)

        super().__init__(
            componentName="image-grid",
            columnSpan=column_span,
            images=images,
            grid_size=grid_size,
            row_captions=row_captions,
            column_captions=column_captions,
            caption=caption,
        )

    def custom_validation(self) -> None:
        num_rows, num_cols = self.data["grid_size"]
        num_images = len(self.data["images"])
        if not (
            num_images == num_cols * num_rows  # Exact number of images
            or (num_images // num_cols == num_rows - 1 and num_images % num_cols != 0)  # Last row incomplete
        ):
            raise ValueError("Grid is too small or too large for the number of images provided.")
        try:
            if len(self.data["row_captions"]) != num_rows:
                raise ValueError("Number of row captions must match the number of rows in the grid.")
        except KeyError:
            # Skip validation if row captions are not provided
            pass
        try:
            if len(self.data["column_captions"]) != num_cols:
                raise ValueError("Number of column captions must match the number of columns in the grid.")
        except KeyError:
            # Skip validation if column captions are not provided
            pass

    def __str__(self) -> str:
        output_dict = deepcopy(self.data)
        output_dict["images"] = [json.loads(str(image)) for image in self._image_objects]
        return json.dumps(output_dict)
