"""YOLO 风格标注的数据集读取。"""

import pytest

pytest.importorskip("torch")
pytest.importorskip("PIL")

from pathlib import Path  # noqa: E402

import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402

from datasets import FilmDefectDataset, collate_fn  # noqa: E402

CLASSES = ["scratch", "bubble"]
TRANSFORMS = [
    {"name": "Resize", "size": [64, 64]},
    {"name": "ToTensor"},
    {"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]},
]


@pytest.fixture
def dataset_root(tmp_path):
    """YOLO 目录：images/ 与同层 labels/，128x128 的图，标注为归一化 cx cy w h。"""
    images = tmp_path / "images"
    labels = tmp_path / "labels"
    images.mkdir()
    labels.mkdir()

    for name in ("a.png", "b.png"):
        Image.fromarray(np.zeros((128, 128, 3), dtype=np.uint8)).save(images / name)
    # 像素框 (10,10,50,50) -> 中心 (30,30)/128, 宽高 40/128
    (labels / "a.txt").write_text("0 0.234375 0.234375 0.312500 0.312500\n", encoding="utf-8")
    (labels / "b.txt").write_text("", encoding="utf-8")  # 空文件 = 无缺陷
    return images, labels


def build(dataset_root, **kwargs):
    images, _ = dataset_root
    return FilmDefectDataset(
        img_dir=images, classes=CLASSES, transforms=TRANSFORMS, stride=4, **kwargs
    )


def test_label_dir_follows_yolo_convention(dataset_root):
    images, labels = dataset_root
    dataset = build(dataset_root)
    assert len(dataset) == 2
    assert dataset.label_file(images / "a.png") == labels / "a.txt"


def test_sample_shapes_and_box_denormalization(dataset_root):
    dataset = build(dataset_root)
    sample = dataset[0]

    assert sample["image"].shape == (3, 64, 64)
    assert sample["target"]["heatmap"].shape == (2, 16, 16)
    # 归一化框 -> 128px 像素框 (10,10,50,50) -> Resize 到 64 后减半
    assert sample["boxes"][0].tolist() == pytest.approx([5.0, 5.0, 25.0, 25.0], abs=1e-3)
    assert sample["labels"].tolist() == [0]
    assert sample["image_label"].tolist() == [1.0, 0.0]


def test_empty_label_file_is_defect_free(dataset_root):
    sample = build(dataset_root)[1]
    assert sample["boxes"].numel() == 0
    assert sample["target"]["mask"].sum() == 0
    assert sample["image_label"].sum() == 0


def test_missing_label_file_is_defect_free(dataset_root):
    _, labels = dataset_root
    (labels / "b.txt").unlink()
    assert build(dataset_root)[1]["boxes"].numel() == 0


def test_explicit_label_dir(dataset_root, tmp_path):
    images, labels = dataset_root
    other = tmp_path / "yolo_labels"
    other.mkdir()
    (other / "a.txt").write_text("1 0.5 0.5 0.25 0.25\n", encoding="utf-8")
    (other / "b.txt").write_text("", encoding="utf-8")

    dataset = build(dataset_root, label_dir=other)
    assert dataset[0]["labels"].tolist() == [1]


def test_list_file_controls_order(dataset_root, tmp_path):
    images, _ = dataset_root
    list_file = tmp_path / "val.txt"
    list_file.write_text("images/b.png\nimages/a.png\n", encoding="utf-8")

    dataset = build(dataset_root, list_file=list_file)
    assert [Path(item).name for item in [dataset[i]["meta"]["path"] for i in range(2)]] == [
        "b.png",
        "a.png",
    ]


def test_collate_keeps_boxes_as_list(dataset_root):
    dataset = build(dataset_root)
    batch = collate_fn([dataset[0], dataset[1]])

    assert batch["images"].shape == (2, 3, 64, 64)
    assert batch["target"]["heatmap"].shape == (2, 2, 16, 16)
    assert isinstance(batch["boxes"], list) and len(batch["boxes"]) == 2
    assert batch["image_labels"].shape == (2, 2)


def test_class_id_out_of_range_raises(dataset_root):
    _, labels = dataset_root
    (labels / "a.txt").write_text("7 0.5 0.5 0.2 0.2\n", encoding="utf-8")
    with pytest.raises(ValueError, match="越界"):
        _ = build(dataset_root)[0]


def test_gray_mode_yields_single_channel(dataset_root):
    images, _ = dataset_root
    transforms = [
        {"name": "Resize", "size": [64, 64]},
        {"name": "ToTensor"},
        {"name": "Normalize", "mean": [0.5], "std": [0.5]},
    ]
    dataset = FilmDefectDataset(
        img_dir=images, classes=CLASSES, transforms=transforms, stride=4, gray=True
    )
    assert dataset[0]["image"].shape == (1, 64, 64)


def test_classes_from_data_yaml(dataset_root, tmp_path):
    pytest.importorskip("yaml")
    from datasets import write_data_yaml

    images, _ = dataset_root
    data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path)
    dataset = FilmDefectDataset(
        img_dir=images, classes=data_yaml, transforms=TRANSFORMS, stride=4
    )
    assert dataset.classes == CLASSES


def test_metrics_on_perfect_prediction():
    from engine.metrics import DetectionMetric

    metric = DetectionMetric(iou_threshold=0.5)
    boxes = torch.tensor([[0.0, 0.0, 10.0, 10.0]])
    metric.update(boxes, torch.tensor([0]), boxes, torch.tensor([0]))
    result = metric.compute()
    assert result["precision"] == 1.0 and result["recall"] == 1.0