from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np

from pyimgano.utils.datasets import MVTecLOCODataset


def _write_rgb(path: Path, *, size=(32, 32), value=128) -> None:
    img = np.ones((size[0], size[1], 3), dtype=np.uint8) * value
    cv2.imwrite(str(path), img)


def _write_mask(path: Path, *, size=(32, 32)) -> None:
    mask = np.zeros((size[0], size[1]), dtype=np.uint8)
    mask[8:24, 8:24] = 255
    cv2.imwrite(str(path), mask)


def test_mvtec_loco_path_accessors(tmp_path: Path) -> None:
    root = tmp_path / "mvtec_loco"
    cat = "breakfast_box"

    (root / cat / "train" / "good").mkdir(parents=True)
    (root / cat / "test" / "good").mkdir(parents=True)

    (root / cat / "test" / "logical_anomalies" / "000").mkdir(parents=True)
    (root / cat / "ground_truth" / "logical_anomalies" / "000").mkdir(parents=True)

    (root / cat / "test" / "structural_anomalies" / "crack").mkdir(parents=True)
    (root / cat / "ground_truth" / "structural_anomalies" / "crack").mkdir(parents=True)

    _write_rgb(root / cat / "train" / "good" / "train_0.png")
    _write_rgb(root / cat / "test" / "good" / "good_0.png")

    _write_rgb(root / cat / "test" / "logical_anomalies" / "000" / "bad_0.png", value=200)
    _write_mask(root / cat / "ground_truth" / "logical_anomalies" / "000" / "bad_0.png")

    _write_rgb(
        root / cat / "test" / "structural_anomalies" / "crack" / "bad_1.png",
        value=220,
    )
    _write_mask(
        root / cat / "ground_truth" / "structural_anomalies" / "crack" / "bad_1_mask.png"
    )

    ds = MVTecLOCODataset(root=str(root), category=cat, resize=(32, 32), load_masks=True)
    train_paths = ds.get_train_paths()
    assert len(train_paths) == 1

    test_paths, labels, masks = ds.get_test_paths()
    assert len(test_paths) == 3
    assert labels.tolist() == [0, 1, 1]
    assert masks is not None
    assert masks.shape == (3, 32, 32)
    assert masks[0].sum() == 0
    assert masks[1].sum() > 0
    assert masks[2].sum() > 0

