"""Tests for zarr_copy.slurm_rechunk."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
import zarr

from zarr_copy.slurm_rechunk import (
    build_job_script,
    build_parser,
    main,
    write_single_var_transform_spec,
)

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


@pytest.fixture
def src_zarr(tmp_path):
    path = str(tmp_path / "src.zarr")
    grp = zarr.open_group(path, mode="w")
    a = grp.create_array("ttr", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f4")
    a[:] = 0
    a.update_attributes({"_ARRAY_DIMENSIONS": ["time", "lat", "lon"]})
    b = grp.create_array("msl", shape=(4, 8, 8), chunks=(2, 4, 4), dtype="f4")
    b[:] = 1
    b.update_attributes({"_ARRAY_DIMENSIONS": ["time", "lat", "lon"]})
    return path


@pytest.fixture
def transform_cfg(tmp_path):
    cfg = tmp_path / "t.py"
    cfg.write_text("transform = {\n    'rlut': lambda ttr: ttr / -3600,\n    'psl': 'msl',\n}\n")
    return cfg


# ---------------------------------------------------------------------------
# build_parser
# ---------------------------------------------------------------------------


class TestBuildParser:
    def test_transform_flag_parsed(self, src_zarr, tmp_path, transform_cfg):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst, "--transform", str(transform_cfg)])
        assert args.transform == transform_cfg

    def test_transform_default_is_none(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        args = build_parser().parse_args([src_zarr, dst])
        assert args.transform is None


# ---------------------------------------------------------------------------
# write_single_var_transform_spec
# ---------------------------------------------------------------------------


class TestWriteSingleVarTransformSpec:
    def test_string_entry(self, tmp_path):
        p = write_single_var_transform_spec("psl", "msl", tmp_path)
        assert p.exists()
        content = p.read_text()
        assert "'psl'" in content
        assert "'msl'" in content

    def test_lambda_entry(self, tmp_path):
        p = write_single_var_transform_spec("rlut", lambda ttr: ttr / -3600, tmp_path)
        assert p.exists()
        content = p.read_text()
        assert "'rlut'" in content
        assert "lambda" in content

    def test_creates_directory(self, tmp_path):
        dest = tmp_path / "subdir" / "specs"
        write_single_var_transform_spec("psl", "msl", dest)
        assert dest.is_dir()

    def test_generated_spec_is_loadable(self, tmp_path):
        """The written file must be importable as a valid transform spec."""
        from zarr_copy.transform import load_transform_spec

        p = write_single_var_transform_spec("psl", "msl", tmp_path)
        spec = load_transform_spec(p)
        assert spec == {"psl": "msl"}

    def test_generated_lambda_spec_is_loadable(self, tmp_path):
        from zarr_copy.transform import load_transform_spec

        p = write_single_var_transform_spec("rlut", lambda ttr: ttr / -3600, tmp_path)
        spec = load_transform_spec(p)
        assert "rlut" in spec
        assert callable(spec["rlut"])
        import numpy as np

        result = spec["rlut"](np.array([-3600.0]))
        assert result[0] == pytest.approx(1.0)


# ---------------------------------------------------------------------------
# build_job_script
# ---------------------------------------------------------------------------


class TestBuildJobScript:
    def _base_kwargs(self, tmp_path):
        return dict(
            varname="ttr",
            src_paths=["/src.zarr"],
            zarr_out="/dst.zarr",
            rechunk=None,
            double_remap=None,
            tmp_path=None,
            cpus=4,
            mem="32G",
            walltime="04:00:00",
            log_dir=tmp_path / "logs",
        )

    def test_no_transform_uses_v_flag(self, tmp_path):
        script = build_job_script(**self._base_kwargs(tmp_path))
        assert "-v ttr" in script
        assert "--transform" not in script

    def test_multiple_sources_in_command(self, tmp_path):
        kw = self._base_kwargs(tmp_path)
        kw["src_paths"] = ["/src_a.zarr", "/src_b.zarr"]
        script = build_job_script(**kw)
        assert "/src_a.zarr" in script
        assert "/src_b.zarr" in script

    def test_transform_uses_transform_flag(self, tmp_path):
        tf = tmp_path / "spec.py"
        kw = self._base_kwargs(tmp_path)
        kw["transform_file"] = tf
        script = build_job_script(**kw)
        assert "--transform" in script
        assert str(tf) in script
        assert "-v ttr" not in script

    def test_rechunk_forwarded(self, tmp_path):
        kw = self._base_kwargs(tmp_path)
        kw["rechunk"] = ["time=1", "lat=256"]
        script = build_job_script(**kw)
        assert "--rechunk time=1 lat=256" in script

    def test_double_remap_forwarded(self, tmp_path):
        kw = self._base_kwargs(tmp_path)
        kw["double_remap"] = "time"
        script = build_job_script(**kw)
        assert "--double-remap time" in script

    def test_sbatch_headers_present(self, tmp_path):
        script = build_job_script(**self._base_kwargs(tmp_path))
        assert "#SBATCH --cpus-per-task=4" in script
        assert "#SBATCH --mem=32G" in script
        assert "#SBATCH --time=04:00:00" in script

    def test_no_consolidated_forwarded(self, tmp_path):
        kw = self._base_kwargs(tmp_path)
        kw["no_consolidated"] = True
        script = build_job_script(**kw)
        assert "--no-consolidated" in script

    def test_no_consolidated_absent_by_default(self, tmp_path):
        script = build_job_script(**self._base_kwargs(tmp_path))
        assert "--no-consolidated" not in script

    def test_verbose_forwarded(self, tmp_path):
        kw = self._base_kwargs(tmp_path)
        kw["verbose"] = True
        script = build_job_script(**kw)
        assert "--verbose" in script

    def test_verbose_absent_by_default(self, tmp_path):
        script = build_job_script(**self._base_kwargs(tmp_path))
        assert "--verbose" not in script

    def test_uses_sys_executable(self, tmp_path):
        import sys

        script = build_job_script(**self._base_kwargs(tmp_path))
        assert sys.executable in script
        assert "-m zarr_copy" in script


# ---------------------------------------------------------------------------
# main() — no-transform path
# ---------------------------------------------------------------------------


class TestMainNoTransform:
    def test_dry_run_no_sbatch(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "-m zarr_copy" in out
        assert "-v" in out

    def test_dry_run_variable_selection(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "-v", "ttr", "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "ttr" in out
        assert "msl" not in out

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

    def test_skips_missing_variable(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "-v", "ttr", "ghost_var", "--dry-run"])
        assert rc == 0
        err = capsys.readouterr().err
        assert "ghost_var" in err

    def test_no_consolidated_forwarded_in_script(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--no-consolidated", "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "--no-consolidated" in out

    def test_no_consolidated_absent_without_flag(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "--no-consolidated" not in out

    def test_verbose_forwarded_in_script(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--verbose", "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "--verbose" in out

    def test_verbose_absent_without_flag(self, src_zarr, tmp_path, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "--verbose" not in out


# ---------------------------------------------------------------------------
# main() — transform path
# ---------------------------------------------------------------------------


class TestMainTransform:
    def test_dry_run_enumerates_output_vars(self, src_zarr, tmp_path, transform_cfg, capsys):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--transform", str(transform_cfg), "--dry-run"])
        assert rc == 0
        out = capsys.readouterr().out
        assert "rlut" in out
        assert "psl" in out

    def test_dry_run_uses_transform_flag(self, src_zarr, tmp_path, transform_cfg, capsys):
        dst = str(tmp_path / "dst.zarr")
        main([src_zarr, dst, "--transform", str(transform_cfg), "--dry-run"])
        out = capsys.readouterr().out
        assert "--transform" in out
        assert "-v rlut" not in out
        assert "-v psl" not in out

    def test_per_var_spec_files_written(self, src_zarr, tmp_path, transform_cfg):
        dst = str(tmp_path / "dst.zarr")
        log_dir = tmp_path / "logs"
        main(
            [
                src_zarr,
                dst,
                "--transform",
                str(transform_cfg),
                "--dry-run",
                "--log-dir",
                str(log_dir),
            ]
        )
        spec_dir = log_dir / "transform_specs"
        assert (spec_dir / "transform_rlut.py").exists()
        assert (spec_dir / "transform_psl.py").exists()

    def test_per_var_spec_files_are_valid(self, src_zarr, tmp_path, transform_cfg):
        from zarr_copy.transform import load_transform_spec

        dst = str(tmp_path / "dst.zarr")
        log_dir = tmp_path / "logs"
        main(
            [
                src_zarr,
                dst,
                "--transform",
                str(transform_cfg),
                "--dry-run",
                "--log-dir",
                str(log_dir),
            ]
        )

        psl_spec = load_transform_spec(log_dir / "transform_specs" / "transform_psl.py")
        assert psl_spec == {"psl": "msl"}

        rlut_spec = load_transform_spec(log_dir / "transform_specs" / "transform_rlut.py")
        assert "rlut" in rlut_spec
        assert callable(rlut_spec["rlut"])

    def test_missing_transform_file_returns_nonzero(self, src_zarr, tmp_path):
        dst = str(tmp_path / "dst.zarr")
        rc = main([src_zarr, dst, "--transform", str(tmp_path / "nope.py"), "--dry-run"])
        assert rc != 0

    def test_rechunk_forwarded_in_script(self, src_zarr, tmp_path, transform_cfg, capsys):
        dst = str(tmp_path / "dst.zarr")
        main([src_zarr, dst, "--transform", str(transform_cfg), "--rechunk", "time=1", "--dry-run"])
        out = capsys.readouterr().out
        assert "--rechunk time=1" in out

    def test_submits_jobs_via_sbatch(self, src_zarr, tmp_path, transform_cfg):
        """Verify sbatch is called once per output variable."""
        dst = str(tmp_path / "dst.zarr")
        mock_result = MagicMock()
        mock_result.stdout = "12345"

        with patch("zarr_copy.slurm_rechunk.subprocess.run", return_value=mock_result) as mock_run:
            rc = main([src_zarr, dst, "--transform", str(transform_cfg)])

        assert rc == 0
        assert mock_run.call_count == 2  # one per output variable in the spec
