"""Orthophoto tiles must reach the model as a faithful 8-bit RGB image.

Belgian orthophoto products are routinely 16-bit and/or 4-band (RGB + NIR).
Handing those to ``PIL.Image.convert("RGB")`` truncates the high byte, so a
bright roof arrives as a near-black pixel and the detector sees nothing that
resembles its training data. The tile is read with rasterio instead, the RGB
bands are selected explicitly and the values are percentile-stretched.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from app.services.yolo_adapter import _prediction_source

rasterio = pytest.importorskip("rasterio")
Image = pytest.importorskip("PIL.Image")


def _write_tile(path: Path, array: np.ndarray, dtype: str) -> None:
    count, height, width = array.shape
    with rasterio.open(
        path,
        "w",
        driver="GTiff",
        width=width,
        height=height,
        count=count,
        dtype=dtype,
    ) as dataset:
        dataset.write(array.astype(dtype))


def _prepared(tile_path: Path) -> np.ndarray:
    with _prediction_source(tile_path) as source:
        with Image.open(source) as image:
            assert image.mode == "RGB"
            return np.array(image)


def test_uint16_tile_keeps_its_contrast_instead_of_going_black(tmp_path: Path) -> None:
    # A typical 12-bit-in-16-bit orthophoto: values well below 65535.
    array = np.zeros((3, 32, 32), dtype=np.uint16)
    array[0] = 800
    array[1] = 1600
    array[2] = 3200
    array[:, 0, 0] = 40  # a dark corner so the stretch has a low anchor
    tile_path = tmp_path / "uint16.tif"
    _write_tile(tile_path, array, "uint16")

    prepared = _prepared(tile_path)

    assert prepared.shape == (32, 32, 3)
    # Naive 16->8 bit truncation would map 800/1600/3200 to near zero.
    assert prepared.max() > 200
    # The three bands stay distinguishable rather than collapsing together.
    assert prepared[16, 16, 0] < prepared[16, 16, 1] < prepared[16, 16, 2]


def test_four_band_rgbi_tile_drops_the_infrared_band(tmp_path: Path) -> None:
    array = np.zeros((4, 16, 16), dtype=np.uint8)
    array[0] = 10
    array[1] = 120
    array[2] = 240
    array[3] = 255  # near-infrared must not be treated as an alpha or a colour
    tile_path = tmp_path / "rgbi.tif"
    _write_tile(tile_path, array, "uint8")

    prepared = _prepared(tile_path)

    assert prepared.shape == (16, 16, 3)
    assert prepared[8, 8, 0] < prepared[8, 8, 1] < prepared[8, 8, 2]


def test_single_band_tile_is_replicated_across_rgb(tmp_path: Path) -> None:
    array = np.full((1, 16, 16), 128, dtype=np.uint8)
    array[0, 0, 0] = 0
    array[0, 15, 15] = 255
    tile_path = tmp_path / "gray.tif"
    _write_tile(tile_path, array, "uint8")

    prepared = _prepared(tile_path)

    assert prepared.shape == (16, 16, 3)
    assert prepared[8, 8, 0] == prepared[8, 8, 1] == prepared[8, 8, 2]


def test_eight_bit_rgb_tile_is_passed_through_unchanged(tmp_path: Path) -> None:
    array = np.zeros((3, 16, 16), dtype=np.uint8)
    array[0] = 10
    array[1] = 120
    array[2] = 240
    tile_path = tmp_path / "rgb.tif"
    _write_tile(tile_path, array, "uint8")

    prepared = _prepared(tile_path)

    # Already display-ready: no stretch should be invented for it.
    assert prepared[8, 8].tolist() == [10, 120, 240]


def test_nodata_pixels_do_not_drive_the_stretch(tmp_path: Path) -> None:
    array = np.full((3, 32, 32), 2000, dtype=np.uint16)
    array[:, :4, :] = 0  # nodata collar from a clipped orthophoto
    tile_path = tmp_path / "nodata.tif"
    with rasterio.open(
        tile_path,
        "w",
        driver="GTiff",
        width=32,
        height=32,
        count=3,
        dtype="uint16",
        nodata=0,
    ) as dataset:
        dataset.write(array)

    prepared = _prepared(tile_path)

    assert prepared[16, 16].tolist() != [0, 0, 0]