"""Tests for zarr_copy.__main__ (CLI)."""

from unittest.mock import MagicMock, call, patch

import numpy as np
import pytest
import zarr

from zarr_copy.__main__ import _open_source_group, build_parser, main


@pytest.fixture
def src_zarr(tmp_path):
    """Create a small source zarr store and return its path."""
    path = str(tmp_path / "src.zarr")
    grp = zarr.open_group(path, mode="a")
    a = grp.create_array("temperature", shape=(6, 6), chunks=(3, 3), dtype="f4")
    a[:] = np.arange(36, dtype="f4").reshape(6, 6)
    a.update_attributes({"_ARRAY_DIMENSIONS": ["x", "y"]})
    b = grp.create_array("salinity", shape=(6, 6), chunks=(3, 3), dtype="f4")
    b[:] = np.zeros((6, 6), dtype="f4")
    b.update_attributes({"_ARRAY_DIMENSIONS": ["x", "y"]})
    return path


class TestOpenSourceGroup:
    def test_regular_path_calls_open_group(self, src_zarr):
        """Regular paths must go through zarr.open_group."""
        with patch("zarr_copy.__main__.zarr.open_group", wraps=zarr.open_group) as mock_open:
            grp = _open_source_group(src_zarr)
        assert isinstance(grp, zarr.Group)
        mock_open.assert_called_once_with(src_zarr, mode="r")

    def test_no_consolidated_forwarded(self, src_zarr):
        with patch("zarr_copy.__main__.zarr.open_group", wraps=zarr.open_group) as mock_open:
            _open_source_group(src_zarr, use_consolidated=False)
        mock_open.assert_called_once_with(src_zarr, mode="r", use_consolidated=False)

    def test_json_path_uses_reference_filesystem(self, tmp_path):
        """Paths ending in .json must open via ReferenceFileSystem + FsspecStore."""
        fake_group = MagicMock(spec=zarr.Group)
        with (
            patch("zarr_copy.__main__.ReferenceFileSystem") as mock_rfs,
            patch("zarr_copy.__main__.FsspecStore") as mock_store,
            patch("zarr_copy.__main__.zarr.open", return_value=fake_group) as mock_open,
        ):
            result = _open_source_group("/some/ref.json")

        mock_rfs.assert_called_once_with("/some/ref.json", asynchronous=True, storage_options={})
        mock_store.assert_called_once_with(fs=mock_rfs.return_value, path="")
        mock_open.assert_called_once_with(store=mock_store.return_value, mode="r", zarr_format=2)
        assert result is fake_group

    def test_json_no_consolidated_forwarded(self, tmp_path):
        """use_consolidated=False must be forwarded to zarr.open for reference paths."""
        fake_group = MagicMock(spec=zarr.Group)
        with (
            patch("zarr_copy.__main__.ReferenceFileSystem"),
            patch("zarr_copy.__main__.FsspecStore"),
            patch("zarr_copy.__main__.zarr.open", return_value=fake_group) as mock_open,
        ):
            _open_source_group("/some/ref.json", use_consolidated=False)
        _, kwargs = mock_open.call_args
        assert kwargs["use_consolidated"] is False

    def test_parq_path_sets_lazy_and_remote_protocol(self, tmp_path):
        """Paths ending in .parq must pass lazy=True and remote_protocol='file'."""
        fake_group = MagicMock(spec=zarr.Group)
        with (
            patch("zarr_copy.__main__.ReferenceFileSystem") as mock_rfs,
            patch("zarr_copy.__main__.FsspecStore"),
            patch("zarr_copy.__main__.zarr.open", return_value=fake_group),
        ):
            _open_source_group("/some/ref.parq")

        _, kwargs = mock_rfs.call_args
        assert kwargs["storage_options"]["lazy"] is True
        assert kwargs["storage_options"]["remote_protocol"] == "file"

    def test_reference_prefix_stripped_from_json(self):
        """reference:: prefix must be stripped before opening a .json reference file."""
        fake_group = MagicMock(spec=zarr.Group)
        with (
            patch("zarr_copy.__main__.ReferenceFileSystem") as mock_rfs,
            patch("zarr_copy.__main__.FsspecStore") as mock_store,
            patch("zarr_copy.__main__.zarr.open", return_value=fake_group) as mock_open,
        ):
            result = _open_source_group("reference::/some/ref.json")

        mock_rfs.assert_called_once_with("/some/ref.json", asynchronous=True, storage_options={})
        mock_store.assert_called_once_with(fs=mock_rfs.return_value, path="")
        mock_open.assert_called_once_with(store=mock_store.return_value, mode="r", zarr_format=2)
        assert result is fake_group

    def test_reference_prefix_stripped_from_parq(self):
        """reference:: prefix must be stripped before opening a .parq reference file."""
        fake_group = MagicMock(spec=zarr.Group)
        with (
            patch("zarr_copy.__main__.ReferenceFileSystem") as mock_rfs,
            patch("zarr_copy.__main__.FsspecStore"),
            patch("zarr_copy.__main__.zarr.open", return_value=fake_group),
        ):
            _open_source_group("reference::/some/ref.parq")

        args, kwargs = mock_rfs.call_args
        assert args[0] == "/some/ref.parq"
        assert kwargs["storage_options"]["lazy"] is True
        assert kwargs["storage_options"]["remote_protocol"] == "file"


class TestBuildParser:
    def test_positional_args(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst])
        assert args.zarr_paths == [src_zarr, dst]

    def test_multiple_sources_parsed(self, src_zarr, tmp_path):
        src2 = str(tmp_path / "src2.zarr")
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, src2, dst])
        assert args.zarr_paths == [src_zarr, src2, dst]

    def test_variable_selection(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst, "-v", "temperature", "salinity"])
        assert args.variables == ["temperature", "salinity"]

    def test_rechunk_parsing(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst, "--rechunk", "x=2", "y=4"])
        assert dict(args.rechunk) == {"x": 2, "y": 4}

    def test_rechunk_bad_format(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        with pytest.raises(SystemExit):
            build_parser().parse_args([src_zarr, dst, "--rechunk", "nodim"])

    def test_rechunk_non_integer_size(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        with pytest.raises(SystemExit):
            build_parser().parse_args([src_zarr, dst, "--rechunk", "x=abc"])

    def test_no_consolidated_default_false(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst])
        assert args.no_consolidated is False

    def test_no_consolidated_flag_parsed(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst, "--no-consolidated"])
        assert args.no_consolidated is True


class TestMain:
    def test_basic_copy(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst])
        assert rc == 0
        grp = zarr.open_group(dst, mode="r")
        assert set(grp.array_keys()) == {"temperature", "salinity"}

    def test_rechunk(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--rechunk", "x=2"])
        assert rc == 0
        grp = zarr.open_group(dst, mode="r")
        assert grp["temperature"].chunks == (2, 3)

    def test_variable_selection(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "-v", "temperature"])
        assert rc == 0
        grp = zarr.open_group(dst, mode="r")
        assert "temperature" in grp
        assert "salinity" not in grp

    def test_data_integrity(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        main([src_zarr, dst, "--rechunk", "x=2"])
        src_grp = zarr.open_group(src_zarr, mode="r")
        dst_grp = zarr.open_group(dst, mode="r")
        np.testing.assert_array_equal(src_grp["temperature"][:], dst_grp["temperature"][:])

    def test_invalid_source_returns_nonzero(self, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main(["/nonexistent/path", dst])
        assert rc != 0

    def test_missing_variable_returns_nonzero(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "-v", "nonexistent"])
        assert rc != 0

    def test_double_remap(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--rechunk", "x=1", "--double-remap", "x"])
        assert rc == 0
        grp = zarr.open_group(dst, mode="r")
        assert grp["temperature"].chunks[0] == 1

    def test_multi_source_without_transform_rejected(self, src_zarr, tmp_path):
        src2 = str(tmp_path / "src2.zarr")
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, src2, dst])
        assert rc != 0

    def test_multi_source_with_transform(self, tmp_path):
        """Variables from two separate datasets combined via a transform."""
        # Dataset A: has 'ttr'
        path_a = str(tmp_path / "a.zarr")
        grp_a = zarr.open_group(path_a, mode="w")
        arr_a = grp_a.create_array("ttr", shape=(4, 4), chunks=(2, 2), dtype="f4")
        arr_a[:] = np.full((4, 4), -3600.0, dtype="f4")
        arr_a.update_attributes({"_ARRAY_DIMENSIONS": ["time", "x"]})

        # Dataset B: has 'msl'
        path_b = str(tmp_path / "b.zarr")
        grp_b = zarr.open_group(path_b, mode="w")
        arr_b = grp_b.create_array("msl", shape=(4, 4), chunks=(2, 2), dtype="f4")
        arr_b[:] = np.full((4, 4), 101325.0, dtype="f4")
        arr_b.update_attributes({"_ARRAY_DIMENSIONS": ["time", "x"]})

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

        rc = main([path_a, path_b, dst, "--transform", str(cfg)])

        assert rc == 0
        out = zarr.open_group(dst, mode="r")
        np.testing.assert_allclose(out["rlut"][:], 1.0, rtol=1e-5)
        np.testing.assert_allclose(out["psl"][:], 101325.0, rtol=1e-5)

    def test_no_consolidated_passed_to_open(self, src_zarr, tmp_path):
        """--no-consolidated must pass use_consolidated=False to zarr.open_group for sources."""
        dst = str(tmp_path / "dst.zarr")
        with patch("zarr_copy.__main__.zarr.open_group", wraps=zarr.open_group) as mock_open:
            rc = main([src_zarr, dst, "--no-consolidated"])
        assert rc == 0
        # The first call opens the source; it must include use_consolidated=False.
        src_call = mock_open.call_args_list[0]
        assert src_call == call(src_zarr, mode="r", use_consolidated=False)

    def test_no_consolidated_not_passed_by_default(self, src_zarr, tmp_path):
        """Without --no-consolidated the keyword must not appear in the source open call."""
        dst = str(tmp_path / "dst.zarr")
        with patch("zarr_copy.__main__.zarr.open_group", wraps=zarr.open_group) as mock_open:
            rc = main([src_zarr, dst])
        assert rc == 0
        src_call = mock_open.call_args_list[0]
        assert "use_consolidated" not in src_call.kwargs

    def test_unexpected_error_logs_traceback(self, src_zarr, tmp_path):
        """Unexpected errors must be logged via logger.exception so the traceback appears."""
        dst = str(tmp_path / "dst.zarr")
        with patch("zarr_copy.__main__.copy_group", side_effect=RuntimeError("boom")):
            with patch("zarr_copy.__main__.logger") as mock_logger:
                rc = main([src_zarr, dst])
        assert rc != 0
        mock_logger.exception.assert_called_once()
