import cv2
import numpy as np
import pytest
from PIL import Image, ImageDraw

from adetailer.mask import (
    bbox_area,
    dilate_erode,
    has_intersection,
    is_all_black,
    mask_invert,
    mask_merge,
    offset,
)


def test_dilate_positive_value():
    img = Image.new("L", (10, 10), color="black")
    draw = ImageDraw.Draw(img)
    draw.rectangle((3, 3, 5, 5), fill="white")
    value = 3

    result = dilate_erode(img, value)

    assert isinstance(result, Image.Image)
    assert result.size == (10, 10)

    expect = np.array(
        [
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        ],
        dtype=np.uint8,
    )
    assert np.array_equal(np.array(result), expect)


def test_offset():
    img = Image.new("L", (10, 10), color="black")
    draw = ImageDraw.Draw(img)
    draw.rectangle((4, 4, 5, 5), fill="white")

    result = offset(img, x=1, y=2)

    assert isinstance(result, Image.Image)
    assert result.size == (10, 10)

    expect = np.array(
        [
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        ],
        dtype=np.uint8,
    )
    assert np.array_equal(np.array(result), expect)


class TestIsAllBlack:
    def test_is_all_black_1(self):
        img = Image.new("L", (10, 10), color="black")
        assert is_all_black(img)

        draw = ImageDraw.Draw(img)
        draw.rectangle((4, 4, 5, 5), fill="white")
        assert not is_all_black(img)

    def test_is_all_black_2(self):
        img = np.zeros((10, 10), dtype=np.uint8)
        assert is_all_black(img)

        img[4:6, 4:6] = 255
        assert not is_all_black(img)

    def test_is_all_black_rgb_image_pil(self):
        img = Image.new("RGB", (10, 10), color="red")
        assert not is_all_black(img)

        img = Image.new("RGBA", (10, 10), color="red")
        assert not is_all_black(img)

    def test_is_all_black_rgb_image_numpy(self):
        img = np.full((10, 10, 4), 127, dtype=np.uint8)
        with pytest.raises(cv2.error):
            is_all_black(img)

        img = np.full((4, 10, 10), 0.5, dtype=np.float32)
        with pytest.raises(cv2.error):
            is_all_black(img)


class TestHasIntersection:
    def test_has_intersection_1(self):
        arr1 = np.array(
            [
                [0, 0, 0, 0],
                [0, 0, 0, 0],
                [0, 0, 0, 0],
                [0, 0, 0, 0],
            ],
            dtype=np.uint8,
        )
        arr2 = arr1.copy()
        assert not has_intersection(arr1, arr2)

    def test_has_intersection_2(self):
        arr1 = np.array(
            [
                [0, 0, 0, 0],
                [0, 255, 255, 0],
                [0, 255, 255, 0],
                [0, 0, 0, 0],
            ],
            dtype=np.uint8,
        )
        arr2 = np.array(
            [
                [0, 0, 0, 0],
                [0, 0, 0, 0],
                [0, 0, 255, 255],
                [0, 0, 255, 255],
            ],
            dtype=np.uint8,
        )
        assert has_intersection(arr1, arr2)

        arr3 = np.array(
            [
                [255, 0, 0, 0],
                [0, 0, 0, 0],
                [0, 0, 0, 255],
                [0, 0, 255, 255],
            ],
            dtype=np.uint8,
        )
        assert not has_intersection(arr1, arr3)

    def test_has_intersection_3(self):
        img1 = Image.new("L", (10, 10), color="black")
        draw1 = ImageDraw.Draw(img1)
        draw1.rectangle((3, 3, 5, 5), fill="white")
        img2 = Image.new("L", (10, 10), color="black")
        draw2 = ImageDraw.Draw(img2)
        draw2.rectangle((6, 6, 8, 8), fill="white")
        assert not has_intersection(img1, img2)

        img3 = Image.new("L", (10, 10), color="black")
        draw3 = ImageDraw.Draw(img3)
        draw3.rectangle((2, 2, 8, 8), fill="white")
        assert has_intersection(img1, img3)

    def test_has_intersection_4(self):
        img1 = Image.new("RGB", (10, 10), color="black")
        draw1 = ImageDraw.Draw(img1)
        draw1.rectangle((3, 3, 5, 5), fill="white")
        img2 = Image.new("RGBA", (10, 10), color="black")
        draw2 = ImageDraw.Draw(img2)
        draw2.rectangle((2, 2, 8, 8), fill="white")
        assert has_intersection(img1, img2)

    def test_has_intersection_5(self):
        img1 = Image.new("RGB", (10, 10), color="black")
        draw1 = ImageDraw.Draw(img1)
        draw1.rectangle((4, 4, 5, 5), fill="white")
        img2 = np.full((10, 10, 4), 255, dtype=np.uint8)
        assert has_intersection(img1, img2)


def test_bbox_area():
    bbox = [0.0, 0.0, 10.0, 10.0]
    assert bbox_area(bbox) == 100


class TestMaskMerge:
    def test_mask_merge(self):
        img1 = Image.new("L", (10, 10), color="black")
        draw1 = ImageDraw.Draw(img1)
        draw1.rectangle((3, 3, 5, 5), fill="white")

        img2 = Image.new("L", (10, 10), color="black")
        draw2 = ImageDraw.Draw(img2)
        draw2.rectangle((6, 6, 8, 8), fill="white")

        merged = mask_merge([img1, img2])
        assert len(merged) == 1

        expect = Image.new("L", (10, 10), color="black")
        draw3 = ImageDraw.Draw(expect)
        draw3.rectangle((3, 3, 5, 5), fill="white")
        draw3.rectangle((6, 6, 8, 8), fill="white")

        assert np.array_equal(np.array(merged[0]), np.array(expect))

    def test_merge_mask_different_size(self):
        img1 = Image.new("L", (10, 10), color="black")
        draw1 = ImageDraw.Draw(img1)
        draw1.rectangle((3, 3, 5, 5), fill="white")

        img2 = Image.new("L", (20, 20), color="black")
        draw2 = ImageDraw.Draw(img2)
        draw2.rectangle((6, 6, 8, 8), fill="white")

        with pytest.raises(
            cv2.error, match="-209:Sizes of input arguments do not match"
        ):
            mask_merge([img1, img2])


def test_mask_invert():
    img = Image.new("L", (10, 10), color="black")
    draw = ImageDraw.Draw(img)
    draw.rectangle((3, 3, 5, 5), fill="white")

    inverted = mask_invert([img])
    assert len(inverted) == 1

    expect = Image.new("L", (10, 10), color="white")
    draw = ImageDraw.Draw(expect)
    draw.rectangle((3, 3, 5, 5), fill="black")

    assert np.array_equal(np.array(inverted[0]), np.array(expect))