"""Tests for zarr_copy.transform and copy_group transform integration."""

from __future__ import annotations

import numpy as np
import pytest
import zarr

from zarr_copy.copy import copy_group, copy_variable_transformed
from zarr_copy.transform import TransformSpec, apply_transform, load_transform_spec, resolve_sources

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def src_group(tmp_path):
    """Source zarr group mimicking a climate dataset."""
    grp = zarr.open_group(str(tmp_path / "src.zarr"), mode="w")
    rng = np.random.default_rng(42)

    def _add(name, data):
        a = grp.create_array(name, shape=data.shape, chunks=(2, 4, 4), dtype="f4")
        a[:] = data
        a.update_attributes({"_ARRAY_DIMENSIONS": ["time", "lat", "lon"]})
        return a

    _add("ttr", rng.random((4, 8, 8)).astype("f4") * -100)  # thermal radiation (neg)
    _add("tisr", rng.random((4, 8, 8)).astype("f4") * 500)  # TOA incident solar
    _add("tsr", rng.random((4, 8, 8)).astype("f4") * 300)  # surface solar
    _add("tp", rng.random((4, 8, 8)).astype("f4") * 0.01)  # total precip (m)
    _add("msl", rng.random((4, 8, 8)).astype("f4") * 10000 + 100000)  # mean sea-level P
    _add("time", np.arange(4, dtype="f8").reshape(4, 1, 1))  # dummy time scalar

    return grp


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


# ---------------------------------------------------------------------------
# resolve_sources
# ---------------------------------------------------------------------------


class TestResolveSources:
    def test_single_param(self):
        assert resolve_sources(lambda ttr: ttr) == ["ttr"]

    def test_two_params(self):
        assert resolve_sources(lambda tisr, tsr: tisr - tsr) == ["tisr", "tsr"]

    def test_no_params(self):
        assert resolve_sources(lambda: 42) == []

    def test_regular_function(self):
        def my_func(a, b, c):
            return a + b + c

        assert resolve_sources(my_func) == ["a", "b", "c"]


# ---------------------------------------------------------------------------
# load_transform_spec
# ---------------------------------------------------------------------------


class TestLoadTransformSpec:
    def test_loads_valid_file(self, tmp_path):
        cfg = tmp_path / "t.py"
        cfg.write_text("transform = {\n    'out': 'in_var',\n    'computed': lambda x: x * 2,\n}\n")
        spec = load_transform_spec(cfg)
        assert "out" in spec
        assert spec["out"] == "in_var"
        assert callable(spec["computed"])

    def test_missing_file_raises(self, tmp_path):
        with pytest.raises(FileNotFoundError):
            load_transform_spec(tmp_path / "nonexistent.py")

    def test_missing_transform_var_raises(self, tmp_path):
        cfg = tmp_path / "bad.py"
        cfg.write_text("x = 1\n")
        with pytest.raises(ValueError, match="must define"):
            load_transform_spec(cfg)

    def test_non_dict_transform_raises(self, tmp_path):
        cfg = tmp_path / "bad2.py"
        cfg.write_text("transform = [1, 2, 3]\n")
        with pytest.raises(ValueError, match="must be a dict"):
            load_transform_spec(cfg)

    def test_accepts_str_path(self, tmp_path):
        cfg = tmp_path / "t.py"
        cfg.write_text("transform = {'a': 'b'}\n")
        spec = load_transform_spec(str(cfg))
        assert spec == {"a": "b"}


# ---------------------------------------------------------------------------
# apply_transform
# ---------------------------------------------------------------------------


class TestApplyTransform:
    def test_rename_copies_data(self, tmp_path, src_group):
        dst_grp = zarr.open_group(str(tmp_path / "dst.zarr"), mode="w")
        dst_var = dst_grp.create_array("psl", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f4")
        region = (slice(0, 4), slice(0, 8), slice(0, 8))

        apply_transform(region, "msl", src_group, dst_var, "psl")

        np.testing.assert_array_equal(src_group["msl"][:], dst_var[:])

    def test_callable_single_source(self, tmp_path, src_group):
        dst_grp = zarr.open_group(str(tmp_path / "dst.zarr"), mode="w")
        dst_var = dst_grp.create_array("rlut", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f8")
        region = (slice(0, 4), slice(0, 8), slice(0, 8))

        apply_transform(region, lambda ttr: ttr / -3600, src_group, dst_var, "rlut")

        expected = src_group["ttr"][:] / -3600
        np.testing.assert_allclose(dst_var[:], expected, rtol=1e-5)

    def test_callable_two_sources(self, tmp_path, src_group):
        dst_grp = zarr.open_group(str(tmp_path / "dst.zarr"), mode="w")
        dst_var = dst_grp.create_array("rsut", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f8")
        region = (slice(0, 4), slice(0, 8), slice(0, 8))

        apply_transform(
            region,
            lambda tisr, tsr: (tisr - tsr) / 3600,
            src_group,
            dst_var,
            "rsut",
        )

        expected = (src_group["tisr"][:] - src_group["tsr"][:]) / 3600
        np.testing.assert_allclose(dst_var[:], expected, rtol=1e-5)

    def test_missing_source_raises(self, tmp_path, src_group):
        dst_grp = zarr.open_group(str(tmp_path / "dst.zarr"), mode="w")
        dst_var = dst_grp.create_array("out", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f4")
        region = (slice(0, 4), slice(0, 8), slice(0, 8))

        with pytest.raises(KeyError, match="no_such_var"):
            apply_transform(region, "no_such_var", src_group, dst_var, "out")

    def test_missing_callable_source_raises(self, tmp_path, src_group):
        dst_grp = zarr.open_group(str(tmp_path / "dst.zarr"), mode="w")
        dst_var = dst_grp.create_array("out", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f8")
        region = (slice(0, 4), slice(0, 8), slice(0, 8))

        with pytest.raises(KeyError, match="ghost"):
            apply_transform(region, lambda ghost: ghost * 2, src_group, dst_var, "out")


# ---------------------------------------------------------------------------
# copy_variable_transformed
# ---------------------------------------------------------------------------


class TestCopyVariableTransformed:
    def test_rename_produces_correct_data(self, tmp_path, src_group, dst_group):
        copy_variable_transformed("psl", "msl", src_group, dst_group)

        assert "psl" in dst_group
        np.testing.assert_array_equal(src_group["msl"][:], dst_group["psl"][:])

    def test_computed_single_source(self, tmp_path, src_group, dst_group):
        copy_variable_transformed("rlut", lambda ttr: ttr / -3600, src_group, dst_group)

        expected = src_group["ttr"][:] / -3600
        np.testing.assert_allclose(dst_group["rlut"][:], expected, rtol=1e-5)

    def test_computed_two_sources(self, tmp_path, src_group, dst_group):
        copy_variable_transformed(
            "rsut", lambda tisr, tsr: (tisr - tsr) / 3600, src_group, dst_group
        )

        expected = (src_group["tisr"][:] - src_group["tsr"][:]) / 3600
        np.testing.assert_allclose(dst_group["rsut"][:], expected, rtol=1e-5)

    def test_rechunk_applied(self, tmp_path, src_group, dst_group):
        copy_variable_transformed("psl", "msl", src_group, dst_group, rechunk={"time": 1})

        assert dst_group["psl"].chunks[0] == 1

    def test_missing_primary_raises(self, tmp_path, src_group, dst_group):
        with pytest.raises(KeyError, match="nope"):
            copy_variable_transformed("out", "nope", src_group, dst_group)

    def test_no_params_callable_raises(self, tmp_path, src_group, dst_group):
        with pytest.raises(ValueError, match="no parameters"):
            copy_variable_transformed("out", lambda: 42, src_group, dst_group)

    def test_attrs_copied_from_primary(self, tmp_path, src_group, dst_group):
        copy_variable_transformed("psl", "msl", src_group, dst_group)
        assert dst_group["psl"].attrs["_ARRAY_DIMENSIONS"] == ["time", "lat", "lon"]


# ---------------------------------------------------------------------------
# copy_group with transform
# ---------------------------------------------------------------------------


class TestCopyGroupTransform:
    def test_full_transform_spec(self, tmp_path, src_group, dst_group):
        spec: TransformSpec = {
            "rlut": lambda ttr: ttr / -3600,
            "rsut": lambda tisr, tsr: (tisr - tsr) / 3600,
            "pr": lambda tp: tp * (1000 / 3600),
            "psl": "msl",
        }

        copy_group(src_group, dst_group, transform=spec)

        assert set(dst_group.array_keys()) == {"rlut", "rsut", "pr", "psl"}
        np.testing.assert_array_equal(dst_group["psl"][:], src_group["msl"][:])
        np.testing.assert_allclose(dst_group["rlut"][:], src_group["ttr"][:] / -3600, rtol=1e-5)
        np.testing.assert_allclose(
            dst_group["rsut"][:],
            (src_group["tisr"][:] - src_group["tsr"][:]) / 3600,
            rtol=1e-5,
        )
        np.testing.assert_allclose(
            dst_group["pr"][:], src_group["tp"][:] * (1000 / 3600), rtol=1e-5
        )

    def test_transform_ignores_variables_arg(self, tmp_path, src_group, dst_group):
        spec: TransformSpec = {"psl": "msl"}
        copy_group(src_group, dst_group, variables=["ttr"], transform=spec)
        # Only 'psl' written; 'ttr' should not appear
        assert "psl" in dst_group
        assert "ttr" not in dst_group

    def test_group_attrs_copied(self, tmp_path, src_group, dst_group):
        src_group.update_attributes({"history": "test"})
        copy_group(src_group, dst_group, transform={"psl": "msl"})
        assert dst_group.attrs["history"] == "test"

    def test_no_transform_still_works(self, tmp_path, src_group, dst_group):
        copy_group(src_group, dst_group, variables=["msl"])
        assert "msl" in dst_group
        np.testing.assert_array_equal(src_group["msl"][:], dst_group["msl"][:])

    def test_transform_with_rechunk(self, tmp_path, src_group, dst_group):
        spec: TransformSpec = {"psl": "msl"}
        copy_group(src_group, dst_group, rechunk={"time": 1}, transform=spec)
        assert dst_group["psl"].chunks[0] == 1


# ---------------------------------------------------------------------------
# CLI integration: --transform flag
# ---------------------------------------------------------------------------


class TestCLITransform:
    def test_transform_flag_end_to_end(self, tmp_path, src_group):
        from zarr_copy.__main__ import main

        cfg = tmp_path / "t.py"
        cfg.write_text(
            "transform = {\n    'psl': 'msl',\n    'rlut': lambda ttr: ttr / -3600,\n}\n"
        )
        src_path = str(tmp_path / "src.zarr")
        dst_path = str(tmp_path / "dst.zarr")

        ret = main([src_path, dst_path, "--transform", str(cfg)])

        assert ret == 0
        dst = zarr.open_group(dst_path, mode="r")
        assert "psl" in dst
        assert "rlut" in dst
        np.testing.assert_array_equal(dst["psl"][:], src_group["msl"][:])

    def test_missing_transform_file_returns_1(self, tmp_path, src_group):
        from zarr_copy.__main__ import main

        src_path = str(tmp_path / "src.zarr")
        dst_path = str(tmp_path / "dst.zarr")

        ret = main([src_path, dst_path, "--transform", str(tmp_path / "nope.py")])
        assert ret == 1
