"""Tests for I/O utilities: save functions, async pool, dtype correctness."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch

from safetensors.torch import load_file as st_load_file

from src.augmentor.io_utils import (
    make_palette,
    npy_path, vis_path, safetensors_path,
    save_chmv2,
    save_chmv2_async,
    save_depth,
    save_depth_async,
    save_edges,
    save_edges_async,
    save_segmentation,
    save_segmentation_async,
    consolidate_safetensors,
    consolidate_safetensors_async,
    shutdown_io_pool,
    _atomic_save_npy,
)


# ---------------------------------------------------------------------------
# Atomic save
# ---------------------------------------------------------------------------

class TestAtomicSaveNpy:
    def test_creates_file(self, tmp_path: Path) -> None:
        arr = np.array([1.0, 2.0, 3.0])
        path = tmp_path / "test.npy"
        _atomic_save_npy(arr, path)
        assert path.exists()
        loaded = np.load(path)
        np.testing.assert_array_equal(loaded, arr)

    def test_overwrites_existing(self, tmp_path: Path) -> None:
        path = tmp_path / "test.npy"
        _atomic_save_npy(np.array([1.0]), path)
        _atomic_save_npy(np.array([2.0]), path)
        loaded = np.load(path)
        assert loaded.item() == 2.0

    def test_no_leftover_tmp_on_success(self, tmp_path: Path) -> None:
        _atomic_save_npy(np.array([1.0]), tmp_path / "test.npy")
        tmp_files = list(tmp_path.glob("*.tmp"))
        assert len(tmp_files) == 0


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

class TestPathHelpers:
    def test_npy_path(self) -> None:
        p = npy_path(Path("/out"), "depth", "Rot/scene/DB/img", "crop_1")
        assert p == Path("/out/npy/depth/Rot/scene/DB/img/crop_1.npy")

    def test_vis_path(self) -> None:
        p = vis_path(Path("/out"), "segm", "Rot/scene/DB/img", "crop_1")
        assert p == Path("/out/segm/Rot/scene/DB/img/crop_1.png")

    def test_safetensors_path(self) -> None:
        p = safetensors_path(Path("/out"), "Rot/scene/DB/img", "crop_1")
        assert p == Path("/out/safetensors/Rot/scene/DB/img/crop_1.safetensors")


# ---------------------------------------------------------------------------
# save_depth
# ---------------------------------------------------------------------------

_RP = "sub"  # rel_parent for tests


class TestSaveDepth:
    def test_saves_float16_npy(self, tmp_path: Path) -> None:
        save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=True, save_vis=False)
        p = npy_path(tmp_path, "depth", _RP, "img01")
        arr = np.load(p)
        assert arr.dtype == np.float16
        assert arr.shape == (1, 32, 32)

    def test_saves_vis_png(self, tmp_path: Path) -> None:
        save_depth(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=False, save_vis=True)
        assert vis_path(tmp_path, "depth", _RP, "img01").exists()
        assert not npy_path(tmp_path, "depth", _RP, "img01").exists()

    def test_npy_values_in_range(self, tmp_path: Path) -> None:
        save_depth(torch.rand(1, 16, 16), tmp_path, _RP, "img01")
        p = npy_path(tmp_path, "depth", _RP, "img01")
        arr = np.load(p).astype(np.float32)
        assert arr.min() >= 0.0
        assert arr.max() <= 1.0


class TestSaveChmv2:
    def test_saves_float16_npy(self, tmp_path: Path) -> None:
        save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=True, save_vis=False)
        arr = np.load(npy_path(tmp_path, "chm", _RP, "img01"))
        assert arr.dtype == np.float16
        assert arr.shape == (1, 32, 32)

    def test_saves_vis_png(self, tmp_path: Path) -> None:
        save_chmv2(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=False, save_vis=True)
        assert vis_path(tmp_path, "chm", _RP, "img01").exists()
        assert not npy_path(tmp_path, "chm", _RP, "img01").exists()

    def test_npy_values_in_range(self, tmp_path: Path) -> None:
        save_chmv2(torch.rand(1, 16, 16), tmp_path, _RP, "img01")
        arr = np.load(npy_path(tmp_path, "chm", _RP, "img01")).astype(np.float32)
        assert arr.min() >= 0.0
        assert arr.max() <= 1.0


class TestSaveEdges:
    def test_saves_float16_npy(self, tmp_path: Path) -> None:
        save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=True, save_vis=False)
        arr = np.load(npy_path(tmp_path, "edge", _RP, "img01"))
        assert arr.dtype == np.float16

    def test_saves_vis(self, tmp_path: Path) -> None:
        save_edges(torch.rand(1, 32, 32), tmp_path, _RP, "img01",
                   save_npy=False, save_vis=True)
        assert vis_path(tmp_path, "edge", _RP, "img01").exists()


class TestSaveSegmentation:
    def test_saves_uint8_npy(self, tmp_path: Path) -> None:
        seg = torch.randint(0, 5, (1, 32, 32), dtype=torch.uint8)
        save_segmentation(seg, tmp_path, _RP, "img01",
                          save_npy=True, save_vis=False, num_classes=5)
        arr = np.load(npy_path(tmp_path, "segm", _RP, "img01"))
        assert arr.dtype == np.uint8
        assert arr.shape == (1, 32, 32)

    def test_saves_vis(self, tmp_path: Path) -> None:
        seg = torch.randint(0, 3, (1, 32, 32), dtype=torch.uint8)
        save_segmentation(seg, tmp_path, _RP, "img01",
                          save_npy=False, save_vis=True, num_classes=3)
        assert vis_path(tmp_path, "segm", _RP, "img01").exists()

    def test_int_tensor_also_works(self, tmp_path: Path) -> None:
        seg = torch.randint(0, 5, (1, 16, 16), dtype=torch.int64)
        save_segmentation(seg, tmp_path, _RP, "img01", num_classes=5)
        arr = np.load(npy_path(tmp_path, "segm", _RP, "img01"))
        assert arr.dtype == np.uint8


class TestAsyncSaves:
    def test_depth_async(self, tmp_path: Path) -> None:
        save_depth_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01",
                         save_npy=True, save_vis=False)
        shutdown_io_pool()
        assert npy_path(tmp_path, "depth", _RP, "img01").exists()

    def test_chmv2_async(self, tmp_path: Path) -> None:
        save_chmv2_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01",
                         save_npy=True, save_vis=False)
        shutdown_io_pool()
        assert npy_path(tmp_path, "chm", _RP, "img01").exists()

    def test_edges_async(self, tmp_path: Path) -> None:
        save_edges_async(torch.rand(1, 16, 16), tmp_path, _RP, "img01",
                         save_npy=True, save_vis=False)
        shutdown_io_pool()
        assert npy_path(tmp_path, "edge", _RP, "img01").exists()

    def test_segmentation_async(self, tmp_path: Path) -> None:
        seg = torch.randint(0, 3, (1, 16, 16), dtype=torch.uint8)
        save_segmentation_async(seg, tmp_path, _RP, "img01",
                                save_npy=True, save_vis=False, num_classes=3)
        shutdown_io_pool()
        assert npy_path(tmp_path, "segm", _RP, "img01").exists()

    def test_multiple_async_writes(self, tmp_path: Path) -> None:
        for i in range(8):
            save_depth_async(torch.rand(1, 8, 8), tmp_path, _RP, f"img_{i}",
                             save_npy=True, save_vis=False)
        shutdown_io_pool()
        for i in range(8):
            assert npy_path(tmp_path, "depth", _RP, f"img_{i}").exists()


# ---------------------------------------------------------------------------
# Palette
# ---------------------------------------------------------------------------

class TestMakePalette:
    def test_shape(self) -> None:
        pal = make_palette(10)
        assert pal.shape == (10, 3)
        assert pal.dtype == np.uint8

    def test_background_is_black(self) -> None:
        pal = make_palette(5)
        np.testing.assert_array_equal(pal[0], [0, 0, 0])

    def test_cache_returns_same(self) -> None:
        pal = make_palette(7)
        pal2 = make_palette(7)
        assert pal is pal2


# ---------------------------------------------------------------------------
# SafeTensors consolidation
# ---------------------------------------------------------------------------

class TestConsolidateSafetensors:
    def _save_all_modalities(self, tmp_path: Path, stem: str = "img01") -> None:
        H, W = 32, 32
        save_depth(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False)
        save_edges(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False)
        save_chmv2(torch.rand(1, H, W), tmp_path, _RP, stem, save_npy=True, save_vis=False)
        save_segmentation(
            torch.randint(0, 11, (1, H, W), dtype=torch.uint8),
            tmp_path, _RP, stem, save_npy=True, save_vis=False, num_classes=11,
        )

    def test_creates_safetensors_file(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        ok = consolidate_safetensors(tmp_path, _RP, "img01")
        assert ok
        assert safetensors_path(tmp_path, _RP, "img01").exists()

    def test_contains_all_modalities(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors(tmp_path, _RP, "img01")
        data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
        assert set(data.keys()) == {"depth", "edge", "chm", "segm"}

    def test_dtypes_correct(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors(tmp_path, _RP, "img01")
        data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
        assert data["depth"].dtype == torch.float16
        assert data["edge"].dtype == torch.float16
        assert data["chm"].dtype == torch.float16
        assert data["segm"].dtype == torch.uint8

    def test_shapes_correct(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors(tmp_path, _RP, "img01")
        data = st_load_file(safetensors_path(tmp_path, _RP, "img01"))
        for key in ("depth", "edge", "chm", "segm"):
            assert data[key].shape == (1, 32, 32), f"{key} shape mismatch"

    def test_partial_modalities(self, tmp_path: Path) -> None:
        save_depth(torch.rand(1, 16, 16), tmp_path, _RP, "img02", save_npy=True, save_vis=False)
        save_edges(torch.rand(1, 16, 16), tmp_path, _RP, "img02", save_npy=True, save_vis=False)
        ok = consolidate_safetensors(tmp_path, _RP, "img02")
        assert ok
        data = st_load_file(safetensors_path(tmp_path, _RP, "img02"))
        assert set(data.keys()) == {"depth", "edge"}

    def test_no_modalities_returns_false(self, tmp_path: Path) -> None:
        ok = consolidate_safetensors(tmp_path, _RP, "missing")
        assert not ok
        assert not safetensors_path(tmp_path, _RP, "missing").exists()

    def test_cleanup_npy(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=True)
        assert safetensors_path(tmp_path, _RP, "img01").exists()
        assert not npy_path(tmp_path, "depth", _RP, "img01").exists()
        assert not npy_path(tmp_path, "edge", _RP, "img01").exists()
        assert not npy_path(tmp_path, "chm", _RP, "img01").exists()
        assert not npy_path(tmp_path, "segm", _RP, "img01").exists()

    def test_no_cleanup_keeps_npy(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors(tmp_path, _RP, "img01", cleanup_npy=False)
        assert safetensors_path(tmp_path, _RP, "img01").exists()
        assert npy_path(tmp_path, "depth", _RP, "img01").exists()

    def test_async_consolidation(self, tmp_path: Path) -> None:
        self._save_all_modalities(tmp_path)
        consolidate_safetensors_async(tmp_path, _RP, "img01")
        shutdown_io_pool()
        assert safetensors_path(tmp_path, _RP, "img01").exists()

    def test_from_png_only(self, tmp_path: Path) -> None:
        H, W = 32, 32
        save_depth(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True)
        save_edges(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True)
        save_chmv2(torch.rand(1, H, W), tmp_path, _RP, "img04", save_npy=False, save_vis=True)
        save_segmentation(
            torch.randint(0, 11, (1, H, W), dtype=torch.uint8),
            tmp_path, _RP, "img04", save_npy=False, save_vis=True, num_classes=11,
        )
        ok = consolidate_safetensors(tmp_path, _RP, "img04")
        assert ok
        data = st_load_file(safetensors_path(tmp_path, _RP, "img04"))
        assert set(data.keys()) == {"depth", "edge", "chm", "segm"}
        assert data["segm"].dtype == torch.uint8

    def test_values_preserved(self, tmp_path: Path) -> None:
        depth = torch.rand(1, 16, 16)
        save_depth(depth, tmp_path, _RP, "img03", save_npy=True, save_vis=False)
        consolidate_safetensors(tmp_path, _RP, "img03")
        data = st_load_file(safetensors_path(tmp_path, _RP, "img03"))
        expected = depth.half()
        torch.testing.assert_close(data["depth"], expected, atol=0, rtol=0)
