"""Tests for dataset discovery, filtering, and AugmentDataset."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from PIL import Image

from src.augmentor.dataset import (
    AugmentDataset,
    ImageRecord,
    attach_output_dirs,
    discover_images,
    filter_completed,
    INCOMPLETE_SCENES,
)
from src.augmentor.io_utils import npy_path


class TestDiscoverImages:
    def test_finds_all_images(self, fake_image_dir: Path) -> None:
        records = discover_images(fake_image_dir)
        assert len(records) == 4

    def test_records_are_sorted(self, fake_image_dir: Path) -> None:
        records = discover_images(fake_image_dir)
        paths = [r.rel_path for r in records]
        assert paths == sorted(paths)

    def test_skips_incomplete_scenes(self, tmp_path: Path) -> None:
        scene = list(INCOMPLETE_SCENES)[0]
        img_dir = tmp_path / scene
        img_dir.mkdir(parents=True)
        Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(img_dir / "a.png")
        records = discover_images(tmp_path)
        assert len(records) == 0

    def test_skips_excluded_extensions(self, tmp_path: Path) -> None:
        (tmp_path / "file.tiff").touch()
        (tmp_path / "file.txt").touch()
        (tmp_path / "good.png").write_bytes(
            Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).tobytes()
        )
        records = discover_images(tmp_path)
        assert all(r.abs_path.suffix in {".png", ".jpg", ".jpeg", ".bmp"} for r in records)

    def test_subset_filter(self, tmp_path: Path) -> None:
        sub = tmp_path / "Rot" / "scene"
        sub.mkdir(parents=True)
        Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(sub / "a.png")
        (tmp_path / "Country").mkdir()
        Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(
            tmp_path / "Country" / "b.png",
        )
        records = discover_images(tmp_path, subset="Rot")
        assert len(records) == 1
        assert "Rot" in records[0].rel_path

    def test_empty_root(self, tmp_path: Path) -> None:
        records = discover_images(tmp_path / "nonexistent")
        assert records == []


class TestAttachOutputDirs:
    def test_output_root_structure(self, fake_image_dir: Path, tmp_path: Path) -> None:
        records = discover_images(fake_image_dir)
        out_root = tmp_path / "output"
        attached = attach_output_dirs(records, out_root)
        for r in attached:
            assert r.output_root == out_root
            assert r.rel_parent == str(Path(r.rel_path).parent)


class TestFilterCompleted:
    def test_all_pending_when_no_output(self, sample_records: list[ImageRecord]) -> None:
        pending, skipped = filter_completed(sample_records, "depth")
        assert len(pending) == len(sample_records)
        assert skipped == 0

    def test_skips_completed(self, sample_records: list[ImageRecord]) -> None:
        r = sample_records[0]
        p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
        p.parent.mkdir(parents=True, exist_ok=True)
        np.save(p, np.zeros((1, 8, 8)))
        pending, skipped = filter_completed(sample_records, "depth")
        assert skipped == 1
        assert len(pending) == len(sample_records) - 1

    def test_unknown_stage_returns_all(self, sample_records: list[ImageRecord]) -> None:
        pending, skipped = filter_completed(sample_records, "unknown_stage")
        assert len(pending) == len(sample_records)
        assert skipped == 0


class TestAugmentDataset:
    def test_len(self, sample_records: list[ImageRecord]) -> None:
        ds = AugmentDataset(sample_records, image_size=32)
        assert len(ds) == len(sample_records)

    def test_getitem_shape(self, sample_records: list[ImageRecord]) -> None:
        ds = AugmentDataset(sample_records, image_size=32)
        item = ds[0]
        assert item["image_raw"].shape == (3, 32, 32)
        assert item["image_raw"].dtype == torch.float32
        assert 0.0 <= item["image_raw"].min()
        assert item["image_raw"].max() <= 1.0

    def test_getitem_keys(self, sample_records: list[ImageRecord]) -> None:
        ds = AugmentDataset(sample_records, image_size=32)
        item = ds[0]
        assert set(item.keys()) == {"image_raw", "rel_path", "stem", "output_root", "rel_parent"}
