"""Unit tests for multi-scale tiling, scale indices, and optional quality gates."""

from __future__ import annotations

import numpy as np

from src.core.services.preprocessing import collect_tiles, extract_tiles_for_inference

TILE_SIZES = [128, 64]


def _full_mask(h: int, w: int) -> np.ndarray:
    return np.full((h, w), 255, dtype=np.uint8)


def test_collect_tiles_emits_both_scales_with_aligned_size() -> None:
    image = np.full((256, 256, 3), 120, dtype=np.uint8)
    mask = _full_mask(256, 256)

    regions = collect_tiles(image, mask, TILE_SIZES, overlap_ratio=0.25, min_mask_coverage=0.8)

    assert regions, "expected tiles for a fully-covered mask"
    scale_ids = {r.scale_idx for r in regions}
    assert scale_ids == {0, 1}
    for r in regions:
        assert r.size == TILE_SIZES[r.scale_idx]
        assert r.coverage == 1.0  # fully inside the all-ones mask


def test_variance_gate_rejects_flat_tiles() -> None:
    # Perfectly uniform image -> grayscale variance 0 -> all tiles rejected when the gate is on.
    image = np.full((256, 256, 3), 200, dtype=np.uint8)
    mask = _full_mask(256, 256)

    kept = collect_tiles(image, mask, TILE_SIZES, 0.25, 0.8)
    assert kept  # baseline: coverage gate alone keeps tiles

    gated = collect_tiles(
        image, mask, TILE_SIZES, 0.25, 0.8,
        variance_gate_enabled=True, min_tile_variance=50.0,
    )
    assert gated == []


def test_highlight_gate_rejects_blown_out_tiles() -> None:
    image = np.full((256, 256, 3), 255, dtype=np.uint8)  # 100% specular highlight
    mask = _full_mask(256, 256)

    gated = collect_tiles(
        image, mask, TILE_SIZES, 0.25, 0.8,
        highlight_gate_enabled=True, highlight_thresh=235, max_highlight_frac=0.40,
    )
    assert gated == []


def test_extract_tiles_for_inference_returns_aligned_arrays() -> None:
    image_bgr = np.full((300, 300, 3), 90, dtype=np.uint8)

    tiles, coords, scale_indices, shape = extract_tiles_for_inference(
        image_bgr,
        tile_sizes=TILE_SIZES,
        overlap_ratio=0.25,
        shrink_factor=0.92,
        min_mask_coverage=0.8,
    )

    assert shape == (300, 300)
    assert len(tiles) == coords.shape[0] == scale_indices.shape[0]
    assert coords.dtype == np.int32
    assert scale_indices.dtype == np.int64
    if len(tiles):
        assert int(scale_indices.max()) < len(TILE_SIZES)
        for tile, scale_idx in zip(tiles, scale_indices):
            assert tile.shape[0] == tile.shape[1] == TILE_SIZES[int(scale_idx)]


def test_extract_tiles_empty_when_no_coverage() -> None:
    # Mask is implicit via ROI; force an impossible coverage so nothing is kept.
    image_bgr = np.full((300, 300, 3), 90, dtype=np.uint8)
    tiles, coords, scale_indices, _ = extract_tiles_for_inference(
        image_bgr,
        tile_sizes=TILE_SIZES,
        overlap_ratio=0.25,
        shrink_factor=0.92,
        min_mask_coverage=1.0001,  # unsatisfiable
    )
    assert tiles == []
    assert coords.shape == (0, 2)
    assert scale_indices.shape == (0,)