"""Tests for zarr_copy.copy."""

import numpy as np
import pytest
import zarr

from zarr_copy.copy import copy_group, copy_variable, simple_remap


@pytest.fixture
def src_group(tmp_path):
    """Source zarr group with two float32 variables."""
    grp = zarr.open_group(str(tmp_path / "src.zarr"), mode="w")
    grp.update_attributes({"history": "original"})

    a = grp.create_array("temperature", shape=(10, 20, 30), chunks=(5, 10, 10), dtype="f4")
    a[:] = np.random.default_rng(0).random((10, 20, 30)).astype("f4")
    a.update_attributes({"_ARRAY_DIMENSIONS": ["time", "lat", "lon"]})

    b = grp.create_array("salinity", shape=(10, 20, 30), chunks=(5, 10, 10), dtype="f4")
    b[:] = np.random.default_rng(1).random((10, 20, 30)).astype("f4")
    b.update_attributes({"_ARRAY_DIMENSIONS": ["time", "lat", "lon"]})

    return grp


@pytest.fixture
def dst_group(tmp_path):
    return zarr.open_group(str(tmp_path / "dst.zarr"), mode="a")


class TestSimpleRemap:
    def test_copies_data(self, tmp_path):
        grp = zarr.open_group(str(tmp_path / "g.zarr"), mode="w")
        src = grp.create_array("src", shape=(4, 4), chunks=(2, 2), dtype="f4")
        dst = grp.create_array("dst", shape=(4, 4), chunks=(2, 2), dtype="f4")
        src[:] = np.arange(16, dtype="f4").reshape(4, 4)

        simple_remap((slice(0, 4), slice(0, 4)), src, dst)

        np.testing.assert_array_equal(src[:], dst[:])


class TestCopyVariable:
    def test_no_rechunk(self, src_group, dst_group):
        copy_variable("temperature", src_group, dst_group)
        assert dst_group["temperature"].chunks == src_group["temperature"].chunks
        np.testing.assert_array_equal(src_group["temperature"][:], dst_group["temperature"][:])

    def test_rechunk_single_dim(self, src_group, dst_group):
        copy_variable("temperature", src_group, dst_group, rechunk={"time": 2})
        assert dst_group["temperature"].chunks == (2, 10, 10)
        np.testing.assert_array_equal(src_group["temperature"][:], dst_group["temperature"][:])

    def test_rechunk_multiple_dims(self, src_group, dst_group):
        copy_variable("temperature", src_group, dst_group, rechunk={"time": 2, "lon": 15})
        assert dst_group["temperature"].chunks == (2, 10, 15)
        np.testing.assert_array_equal(src_group["temperature"][:], dst_group["temperature"][:])

    def test_rechunk_unknown_dim_ignored(self, src_group, dst_group):
        copy_variable("temperature", src_group, dst_group, rechunk={"depth": 3})
        # Unknown dim -> chunks unchanged
        assert dst_group["temperature"].chunks == src_group["temperature"].chunks

    def test_attributes_copied(self, src_group, dst_group):
        copy_variable("temperature", src_group, dst_group)
        assert dst_group["temperature"].attrs["_ARRAY_DIMENSIONS"] == [
            "time",
            "lat",
            "lon",
        ]

    def test_double_remap_dim(self, src_group, dst_group):
        copy_variable(
            "temperature",
            src_group,
            dst_group,
            rechunk={"time": 1},
            double_remap_dim="time",
        )
        assert dst_group["temperature"].chunks == (1, 10, 10)
        np.testing.assert_array_equal(src_group["temperature"][:], dst_group["temperature"][:])


class TestCopyGroup:
    def test_copies_all_variables(self, src_group, dst_group):
        copy_group(src_group, dst_group)
        assert set(dst_group.array_keys()) == {"temperature", "salinity"}

    def test_copies_group_attributes(self, src_group, dst_group):
        copy_group(src_group, dst_group)
        assert dst_group.attrs["history"] == "original"

    def test_variable_selection(self, src_group, dst_group):
        copy_group(src_group, dst_group, variables=["temperature"])
        assert "temperature" in dst_group
        assert "salinity" not in dst_group

    def test_rechunk_applied_to_all(self, src_group, dst_group):
        copy_group(src_group, dst_group, rechunk={"time": 2})
        assert dst_group["temperature"].chunks[0] == 2
        assert dst_group["salinity"].chunks[0] == 2

    def test_missing_variable_raises(self, src_group, dst_group):
        with pytest.raises(KeyError, match="nonexistent"):
            copy_group(src_group, dst_group, variables=["nonexistent"])

    def test_data_integrity(self, src_group, dst_group):
        copy_group(src_group, dst_group, rechunk={"lat": 5, "lon": 15})
        for name in ("temperature", "salinity"):
            np.testing.assert_array_equal(src_group[name][:], dst_group[name][:])
