scan4llm/tests/test_slicer.py

from __future__ import annotations

import tempfile
from pathlib import Path

import cv2
import numpy as np
import pytest

from src.split.slicer import (
    SUPPORTED_EXTENSIONS,
    VALID_ROTATIONS,
    add_border,
    crop_to_content,
    detect_grid,
    find_content_bounds,
    generate_output_paths,
    load_image,
    parse_slice,
    process_image,
    rotate_image,
    save_page,
    slice_grid,
)


def _create_test_image(width: int = 400, height: int = 300) -> np.ndarray:
    image = np.zeros((height, width, 3), dtype=np.uint8)
    image[:, :] = (200, 200, 200)
    return image


def _image_to_path(image: np.ndarray, suffix: str = ".png") -> Path:
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        cv2.imwrite(tmp.name, image)
        return Path(tmp.name)


class TestParseSlice:
    def test_simple_2x2(self) -> None:
        assert parse_slice("2:2") == (2, 2)

    def test_2x4(self) -> None:
        assert parse_slice("2:4") == (2, 4)

    def test_1x1(self) -> None:
        assert parse_slice("1:1") == (1, 1)

    def test_large(self) -> None:
        assert parse_slice("10:10") == (10, 10)

    def test_invalid_format_missing_colon(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("22")

    def test_invalid_format_extra_colon(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("2:3:4")

    def test_non_integer_value(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("a:b")

    def test_non_integer_col(self) -> None:
        with pytest.raises(ValueError, match="Неверный формат --slice"):
            parse_slice("x:2")

    def test_zero_cols(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("0:2")

    def test_zero_rows(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("2:0")

    def test_negative_cols(self) -> None:
        with pytest.raises(ValueError, match="Неверное значение --slice"):
            parse_slice("-1:2")

    def test_blank_string(self) -> None:
        with pytest.raises(ValueError):
            parse_slice("")


class TestSliceGrid:
    def test_2x2_exact(self) -> None:
        image = _create_test_image(400, 300)
        pages = slice_grid(image, 2, 2)
        assert len(pages) == 4
        assert pages[0].shape == (150, 200, 3)
        assert pages[1].shape == (150, 200, 3)
        assert pages[2].shape == (150, 200, 3)
        assert pages[3].shape == (150, 200, 3)

    def test_1x1(self) -> None:
        image = _create_test_image(100, 100)
        pages = slice_grid(image, 1, 1)
        assert len(pages) == 1
        assert pages[0].shape == (100, 100, 3)

    def test_2x3(self) -> None:
        image = _create_test_image(300, 200)
        pages = slice_grid(image, 2, 3)
        assert len(pages) == 6
        for page in pages:
            assert page.shape == (100, 100, 3)

    def test_uneven_division_width(self) -> None:
        image = _create_test_image(width=403, height=300)
        pages = slice_grid(image, 1, 2)
        assert len(pages) == 2
        assert pages[0].shape == (300, 201, 3)
        assert pages[1].shape == (300, 202, 3)

    def test_uneven_division_height(self) -> None:
        image = _create_test_image(width=400, height=301)
        pages = slice_grid(image, 2, 1)
        assert len(pages) == 2
        assert pages[0].shape == (150, 400, 3)
        assert pages[1].shape == (151, 400, 3)

    def test_uneven_both(self) -> None:
        image = _create_test_image(width=401, height=301)
        pages = slice_grid(image, 2, 2)
        assert len(pages) == 4
        assert pages[0].shape == (150, 200, 3)
        assert pages[1].shape == (150, 201, 3)
        assert pages[2].shape == (151, 200, 3)
        assert pages[3].shape == (151, 201, 3)

    def test_pixel_content_preserved(self) -> None:
        image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
        pages = slice_grid(image, 2, 2)
        assert np.array_equal(image[0:50, 0:50], pages[0])
        assert np.array_equal(image[0:50, 50:100], pages[1])
        assert np.array_equal(image[50:100, 0:50], pages[2])
        assert np.array_equal(image[50:100, 50:100], pages[3])

    def test_total_pixels_preserved(self) -> None:
        image = _create_test_image(width=401, height=301)
        pages = slice_grid(image, 3, 3)
        total = sum(page.shape[0] * page.shape[1] for page in pages)
        assert total == 401 * 301

    def test_order_2x2(self) -> None:
        image = _create_test_image(200, 200)
        image[0:100, 0:100] = (255, 0, 0)
        image[0:100, 100:200] = (0, 255, 0)
        image[100:200, 0:100] = (0, 0, 255)
        image[100:200, 100:200] = (255, 255, 0)
        pages = slice_grid(image, 2, 2)
        assert pages[0][0, 0].tolist() == [255, 0, 0]
        assert pages[1][0, 0].tolist() == [0, 255, 0]
        assert pages[2][0, 0].tolist() == [0, 0, 255]
        assert pages[3][0, 0].tolist() == [255, 255, 0]

    def test_order_2x4(self) -> None:
        image = _create_test_image(400, 200)
        pages = slice_grid(image, 2, 4)
        assert len(pages) == 8
        assert pages[0].shape == (100, 100, 3)
        assert pages[7].shape == (100, 100, 3)

    def test_4x4(self) -> None:
        image = _create_test_image(400, 400)
        pages = slice_grid(image, 4, 4)
        assert len(pages) == 16
        for page in pages:
            assert page.shape == (100, 100, 3)


class TestAddBorder:
    def test_default_border(self) -> None:
        image = _create_test_image(100, 100)
        result = add_border(image, 20)
        assert result.shape == (140, 140, 3)

    def test_zero_border(self) -> None:
        image = _create_test_image(100, 100)
        result = add_border(image, 0)
        assert np.array_equal(result, image)

    def test_border_is_white(self) -> None:
        image = _create_test_image(100, 100)
        border = 10
        result = add_border(image, border)
        assert np.all(result[0:border, :] == 255)
        assert np.all(result[-border:, :] == 255)
        assert np.all(result[:, 0:border] == 255)
        assert np.all(result[:, -border:] == 255)

    def test_large_border(self) -> None:
        image = _create_test_image(10, 10)
        result = add_border(image, 50)
        assert result.shape == (110, 110, 3)


class TestRotateImage:
    def test_rotate_90(self) -> None:
        image = _create_test_image(200, 100)
        result = rotate_image(image, 90)
        assert result.shape == (200, 100, 3)

    def test_rotate_180(self) -> None:
        image = _create_test_image(200, 100)
        result = rotate_image(image, 180)
        assert result.shape == (100, 200, 3)

    def test_rotate_270(self) -> None:
        image = _create_test_image(200, 100)
        result = rotate_image(image, 270)
        assert result.shape == (200, 100, 3)

    def test_invalid_angle(self) -> None:
        image = _create_test_image()
        with pytest.raises(ValueError, match="Недопустимый угол поворота"):
            rotate_image(image, 45)

    def test_invalid_angle_zero(self) -> None:
        image = _create_test_image()
        with pytest.raises(ValueError):
            rotate_image(image, 0)

    def test_valid_rotations_set(self) -> None:
        assert VALID_ROTATIONS == frozenset({90, 180, 270})


class TestGenerateOutputPaths:
    def test_simple(self) -> None:
        paths = generate_output_paths(Path("scan001.jpg"), 2, Path("/tmp"))
        assert paths == [Path("/tmp/scan001_01.jpg"), Path("/tmp/scan001_02.jpg")]

    def test_tiff_extension(self) -> None:
        paths = generate_output_paths(Path("doc.tiff"), 3, Path("out"))
        assert paths == [
            Path("out/doc_01.tiff"),
            Path("out/doc_02.tiff"),
            Path("out/doc_03.tiff"),
        ]

    def test_png_extension(self) -> None:
        paths = generate_output_paths(Path("img.PNG"), 1, Path("out"))
        assert paths == [Path("out/img_01.png")]

    def test_zero_padding(self) -> None:
        paths = generate_output_paths(Path("scan.jpg"), 12, Path("out"))
        assert paths[0] == Path("out/scan_01.jpg")
        assert paths[9] == Path("out/scan_10.jpg")
        assert paths[11] == Path("out/scan_12.jpg")

    def test_many_pages(self) -> None:
        paths = generate_output_paths(Path("scan.jpg"), 100, Path("out"))
        assert paths[0] == Path("out/scan_001.jpg")
        assert paths[99] == Path("out/scan_100.jpg")

    def test_dot_in_filename(self) -> None:
        paths = generate_output_paths(Path("scan.001.jpg"), 2, Path("out"))
        assert paths == [Path("out/scan.001_01.jpg"), Path("out/scan.001_02.jpg")]


class TestLoadImage:
    def test_valid_jpeg(self) -> None:
        image = _create_test_image()
        path = _image_to_path(image, ".jpg")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_valid_png(self) -> None:
        image = _create_test_image()
        path = _image_to_path(image, ".png")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_valid_tiff(self) -> None:
        image = _create_test_image()
        path = _image_to_path(image, ".tiff")
        loaded = load_image(path)
        assert loaded.shape == (300, 400, 3)

    def test_file_not_found(self) -> None:
        with pytest.raises(FileNotFoundError, match="Файл не найден"):
            load_image(Path("/tmp/nonexistent_scan2html_test.jpg"))

    def test_unsupported_extension(self) -> None:
        path = Path("/tmp/test.bmp")
        with pytest.raises(ValueError, match="Неподдерживаемый формат"):
            load_image(path)

    def test_unsupported_extension_no_dot(self) -> None:
        path = Path("/tmp/test")
        with pytest.raises(ValueError, match="Неподдерживаемый формат"):
            load_image(path)

    def test_supported_extensions_set(self) -> None:
        assert SUPPORTED_EXTENSIONS == frozenset({".jpg", ".jpeg", ".png", ".tiff", ".tif"})


class TestSavePage:
    def test_saves_file(self) -> None:
        image = _create_test_image()
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "test.png"
            save_page(image, path)
            assert path.exists()
            loaded = cv2.imread(str(path))
            assert loaded is not None
            assert loaded.shape == (300, 400, 3)

    def test_creates_parent_dir(self) -> None:
        image = _create_test_image()
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "subdir" / "test.png"
            save_page(image, path)
            assert path.exists()


class TestProcessImage:
    def test_end_to_end_2x2(self) -> None:
        image = _create_test_image(400, 300)
        path = _image_to_path(image, ".jpg")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp))
            assert len(result) == 4
            for p in result:
                assert p.exists()
                loaded = cv2.imread(str(p))
                assert loaded.shape == (250, 300, 3)

    def test_end_to_end_with_border(self) -> None:
        image = _create_test_image(200, 200)
        path = _image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp), border_px=10)
            assert len(result) == 4
            for p in result:
                assert p.exists()
                loaded = cv2.imread(str(p))
                assert loaded.shape == (120, 120, 3)

    def test_end_to_end_tiff(self) -> None:
        image = _create_test_image(400, 400)
        path = _image_to_path(image, ".tiff")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp))
            assert len(result) == 4
            for p in result:
                assert p.suffix.lower() == ".tiff"
                assert p.exists()

    def test_end_to_end_with_rotation(self) -> None:
        image = _create_test_image(200, 400)
        path = _image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 2, 2, Path(tmp), pre_rotate=90)
            assert len(result) == 4
            for p in result:
                assert p.exists()
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape[0] > 0


class TestFindContentBounds:
    def test_uniform_image(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert (x1, y1, x2, y2) == (0, 0, 200, 150)

    def test_single_dot_center(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[75, 100] = (0, 0, 0)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert abs(x1 - 100) <= 1 and abs(x2 - 100) <= 1
        assert abs(y1 - 75) <= 1 and abs(y2 - 75) <= 1

    def test_rectangle_content(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[30:100, 50:140] = (0, 0, 0)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert x1 == 50
        assert y1 == 30
        assert x2 == 139
        assert y2 == 99

    def test_content_at_edges(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[0:10, :] = (0, 0, 0)
        image[:, 0:10] = (0, 0, 0)
        x1, y1, _x2, _y2 = find_content_bounds(image)
        assert y1 == 0
        assert x1 == 0

    def test_grayscale_content(self) -> None:
        image = _create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[20:30, 20:30] = (128, 128, 128)
        x1, y1, x2, y2 = find_content_bounds(image)
        assert x1 <= 29 and x2 >= 20
        assert y1 <= 29 and y2 >= 20


class TestCropToContent:
    def test_crop_centered_content(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[30:50, 40:60] = (0, 0, 0)
        result = crop_to_content(image, border_px=10)
        assert result.shape[0] >= 20
        assert result.shape[1] >= 20

    def test_crop_uniform_image_no_change(self) -> None:
        image = _create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        result = crop_to_content(image, border_px=10)
        assert result.shape == (100, 100, 3)

    def test_border_capped_by_image_edge(self) -> None:
        image = _create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[10:30, 10:30] = (0, 0, 0)
        result = crop_to_content(image, border_px=200)
        assert result.shape == (100, 100, 3)

    def test_border_equals_distance_to_edge(self) -> None:
        image = _create_test_image(100, 100)
        image[:, :] = (255, 255, 255)
        image[40:60, 30:70] = (0, 0, 0)
        result = crop_to_content(image, border_px=50)
        assert result.shape[1] <= 100
        assert result.shape[0] <= 100

    def test_dark_bg_no_crop_uniform(self) -> None:
        image = _create_test_image(100, 100)
        image[:, :] = (0, 0, 0)
        result = crop_to_content(image, border_px=0)
        assert result.shape == (100, 100, 3)

    def test_zero_border_exact_crop(self) -> None:
        image = _create_test_image(200, 150)
        image[:, :] = (255, 255, 255)
        image[20:100, 30:170] = (0, 0, 0)
        result = crop_to_content(image, border_px=0)
        assert result.shape == (80, 140, 3)


class TestProcessImagePostCrop:
    def test_post_crop_reduces_size(self) -> None:
        image = _create_test_image(400, 400)
        image[:, :] = (255, 255, 255)
        image[100:200, 100:200] = (0, 0, 0)
        path = _image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 1, 1, Path(tmp), border_px=30, post_crop=True)
            assert len(result) == 1
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape[0] < 400 or loaded.shape[1] < 400

    def test_post_crop_off_no_crop(self) -> None:
        image = _create_test_image(400, 300)
        path = _image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result = process_image(Path(path), 1, 1, Path(tmp), post_crop=False)
            assert len(result) == 1
            loaded = cv2.imread(str(result[0]))
            assert loaded.shape == (400, 500, 3)

    def test_post_crop_e2e_2x2(self) -> None:
        image = _create_test_image(400, 400)
        image[:, :] = (255, 255, 255)
        image[50:150, 50:150] = (0, 0, 0)
        image[50:150, 250:350] = (0, 0, 0)
        image[250:350, 50:150] = (0, 0, 0)
        image[250:350, 250:350] = (0, 0, 0)
        path = _image_to_path(image, ".png")
        with tempfile.TemporaryDirectory() as tmp:
            result_crop = process_image(Path(path), 2, 2, Path(tmp) / "crop", border_px=30, post_crop=True)
            result_no = process_image(Path(path), 2, 2, Path(tmp) / "no", border_px=30, post_crop=False)
            assert len(result_crop) == 4
            for cp, np in zip(result_crop, result_no):
                c = cv2.imread(str(cp))
                n = cv2.imread(str(np))
                assert c.shape[0] < n.shape[0]
                assert c.shape[1] < n.shape[1]


class TestDetectGrid:
    def test_2x2(self) -> None:
        image = _create_test_image(400, 300)
        image[:, :] = (255, 255, 255)
        dark = (0, 0, 0)
        gap = 20
        bw = (400 - 3 * gap) // 2
        bh = (300 - 3 * gap) // 2
        for r in range(2):
            for c in range(2):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y:y + bh, x:x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 2
        assert rows == 2

    def test_2x4(self) -> None:
        image = _create_test_image(800, 400)
        image[:, :] = (255, 255, 255)
        gap = 15
        bw = (800 - 5 * gap) // 4
        bh = (400 - 3 * gap) // 2
        dark = (0, 0, 0)
        for r in range(2):
            for c in range(4):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y:y + bh, x:x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 4
        assert rows == 2

    def test_1x1(self) -> None:
        image = _create_test_image(200, 200)
        image[:, :] = (255, 255, 255)
        image[20:180, 20:180] = (0, 0, 0)
        cols, rows = detect_grid(image)
        assert cols == 1
        assert rows == 1

    def test_3x3(self) -> None:
        image = _create_test_image(600, 450)
        image[:, :] = (255, 255, 255)
        dark = (0, 0, 0)
        gap = 10
        bw = (600 - 4 * gap) // 3
        bh = (450 - 4 * gap) // 3
        for r in range(3):
            for c in range(3):
                y = gap + r * (bh + gap)
                x = gap + c * (bw + gap)
                image[y:y + bh, x:x + bw] = dark
        cols, rows = detect_grid(image)
        assert cols == 3
        assert rows == 3

    def test_uniform_image(self) -> None:
        image = _create_test_image()
        cols, rows = detect_grid(image)
        assert cols == 1
        assert rows == 1