from unittest.mock import Mock

import torch
from PIL import Image

from utils.image import segment_image, segment_image_dataset


def test_segment_image_passes_pil_image_to_mask_generator() -> None:
    mock_generator = Mock(return_value={"masks": []})

    segment_image(
        mock_generator,
        Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
        points_per_batch=32,
    )

    image_arg = mock_generator.call_args.args[0]
    assert isinstance(image_arg, Image.Image)
    assert image_arg.mode == "RGB"
    assert mock_generator.call_args.kwargs["points_per_batch"] == 32


def test_segment_image_supports_tensor_masks_output() -> None:
    masks_tensor = torch.tensor(
        [
            [
                [1, 1, 0],
                [1, 1, 0],
                [0, 0, 0],
            ],
            [
                [1, 1, 1],
                [1, 1, 1],
                [1, 1, 1],
            ],
        ],
        dtype=torch.float32,
    )
    mock_generator = Mock(return_value={"masks": masks_tensor})

    result = segment_image(
        mock_generator,
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        min_area=3,
        max_masks=5,
    )

    assert len(result) == 2
    assert result[0]["area"] == 9
    assert result[0]["bbox"] == [0, 0, 3, 3]
    assert result[1]["area"] == 4
    assert result[1]["bbox"] == [0, 0, 2, 2]


def test_segment_image_filters_tensor_masks_by_min_area() -> None:
    masks_tensor = torch.tensor(
        [
            [
                [1, 0, 0],
                [0, 0, 0],
                [0, 0, 0],
            ],
            [
                [1, 1, 0],
                [1, 1, 0],
                [0, 0, 0],
            ],
        ],
        dtype=torch.float32,
    )
    mock_generator = Mock(return_value={"masks": masks_tensor})

    result = segment_image(
        mock_generator,
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        min_area=2,
        max_masks=5,
    )

    assert len(result) == 1
    assert result[0]["area"] == 4


def test_segment_image_dataset_returns_per_image_masks_in_order() -> None:
    first_masks = {
        "masks": torch.tensor(
            [[[1, 1, 0], [1, 1, 0], [0, 0, 0]]],
            dtype=torch.float32,
        )
    }
    second_masks = {
        "masks": torch.tensor(
            [[[1, 1, 1], [1, 1, 1], [1, 1, 1]]],
            dtype=torch.float32,
        )
    }
    mock_generator = Mock(side_effect=[first_masks, second_masks])
    images = [
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
        Image.new("RGB", (3, 3), color=(0, 0, 0)),
    ]

    result = segment_image_dataset(
        mock_generator,
        images,
        min_area=2,
        max_masks=5,
        points_per_batch=16,
    )

    assert len(result) == 2
    assert result[0][0]["area"] == 4
    assert result[1][0]["area"] == 9
    assert mock_generator.call_count == 2