import multiprocessing
import os
import signal
import tempfile
from pathlib import Path

import pyarrow.fs
import pytest
import torch

import ray
from ray.tests.client_test_utils import create_remote_signal_actor
from ray.train import BackendConfig, Checkpoint, RunConfig, ScalingConfig, UserCallback
from ray.train.backend import Backend
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR, _get_ray_train_session_dir
from ray.train.tests.util import create_dict_checkpoint
from ray.train.v2._internal.constants import is_v2_enabled
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
from ray.train.v2.api.exceptions import TrainingFailedError, WorkerGroupError
from ray.train.v2.api.result import Result

assert is_v2_enabled()


@pytest.fixture(scope="module", autouse=True)
def ray_start_4_cpus():
    ray.init(num_cpus=4)
    yield
    ray.shutdown()


def test_backend_setup(tmp_path):
    class ValidationBackend(Backend):
        def on_start(self, worker_group, backend_config):
            tmp_path.joinpath("on_start").touch()

        def on_training_start(self, worker_group, backend_config):
            tmp_path.joinpath("on_training_start").touch()

        def on_shutdown(self, worker_group, backend_config):
            tmp_path.joinpath("on_shutdown").touch()

    class ValidationBackendConfig(BackendConfig):
        @property
        def backend_cls(self):
            return ValidationBackend

    trainer = DataParallelTrainer(
        lambda: None,
        backend_config=ValidationBackendConfig(),
        scaling_config=ScalingConfig(num_workers=2),
    )
    trainer.fit()

    assert tmp_path.joinpath("on_start").exists()
    assert tmp_path.joinpath("on_training_start").exists()
    assert tmp_path.joinpath("on_shutdown").exists()


def test_result_output(tmp_path):
    trainer = DataParallelTrainer(
        lambda: None,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(name="test", storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert isinstance(result, Result)
    assert result.path == str(tmp_path / "test")
    assert isinstance(result.filesystem, pyarrow.fs.FileSystem)


def test_no_optional_arguments():
    """Check that the DataParallelTrainer can be instantiated without optional arguments."""
    trainer = DataParallelTrainer(lambda: "not used")
    trainer.fit()


def test_train_loop_config():
    """Check that the train loop config is passed to the train function
    if a config parameter is accepted."""

    def train_fn(config):
        with create_dict_checkpoint({}) as checkpoint:
            ray.train.report(metrics=config, checkpoint=checkpoint)

    train_loop_config = {"x": 1}
    trainer = DataParallelTrainer(
        train_fn,
        train_loop_config=train_loop_config,
        scaling_config=ScalingConfig(num_workers=2),
    )
    result = trainer.fit()
    assert result.metrics == train_loop_config


def test_report_checkpoint_rank0(tmp_path):
    """Check that checkpoints can be reported from rank 0 only."""

    def train_fn():
        metrics = {"rank": ray.train.get_context().get_world_rank()}
        if ray.train.get_context().get_world_rank() == 0:
            with create_dict_checkpoint({}) as checkpoint:
                ray.train.report(metrics=metrics, checkpoint=checkpoint)
        else:
            ray.train.report(metrics=metrics, checkpoint=None)

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert result.metrics == {"rank": 0}
    assert result.checkpoint


def test_report_checkpoint_multirank(tmp_path):
    """Check that checkpoints can be reported from multiple ranks."""

    ranks_to_report = [1, 3]

    def train_fn():
        rank = ray.train.get_context().get_world_rank()
        metrics = {"rank": rank}
        if rank in ranks_to_report:
            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
                Path(temp_checkpoint_dir).joinpath(str(rank)).touch()
                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
                ray.train.report(metrics=metrics, checkpoint=checkpoint)
        else:
            ray.train.report(metrics=metrics, checkpoint=None)

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=4),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert result.checkpoint
    result.checkpoint.to_directory(tmp_path / "validate")
    for rank in ranks_to_report:
        assert tmp_path.joinpath("validate", str(rank)).exists()


def test_error(tmp_path):
    def _error_func_rank_0():
        """An example train_fun that raises an error on rank 0."""
        if ray.train.get_context().get_world_rank() == 0:
            raise ValueError("user error")

    trainer = DataParallelTrainer(
        _error_func_rank_0,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(name="test", storage_path=str(tmp_path)),
    )
    with pytest.raises(TrainingFailedError) as exc_info:
        trainer.fit()
    assert isinstance(exc_info.value, WorkerGroupError)
    assert "user error" in str(exc_info.value.worker_failures[0])
    assert len(exc_info.value.worker_failures) == 1


@pytest.mark.parametrize("env_disabled", [True, False])
def test_setup_working_directory(tmp_path, monkeypatch, env_disabled):
    # Set the environment variable to control the working directory setup
    monkeypatch.setenv(RAY_CHDIR_TO_TRIAL_DIR, str(int(not env_disabled)))

    experiment_dir_name = "test"
    reference_working_dir = (
        Path(_get_ray_train_session_dir(), "test").resolve().as_posix()
    )

    def _check_same_working_directory():
        worker_working_dir = os.getcwd()
        if env_disabled:
            assert worker_working_dir != reference_working_dir
        else:
            assert worker_working_dir == reference_working_dir

    trainer = DataParallelTrainer(
        _check_same_working_directory,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(name=experiment_dir_name, storage_path=str(tmp_path)),
    )
    trainer.fit()


def test_user_callback(tmp_path):
    """Test end to end usage of user callbacks."""
    num_workers = 2

    class MyUserCallback(UserCallback):
        def after_report(self, run_context, metrics, checkpoint):
            assert len(metrics) == num_workers
            assert not checkpoint

        def after_exception(self, run_context, worker_exceptions):
            assert len(worker_exceptions) == 1
            assert worker_exceptions.get(0) is not None

    def _train_fn(config):
        ray.train.report(metrics={"rank": ray.train.get_context().get_world_rank()})
        if ray.train.get_context().get_world_rank() == 0:
            raise ValueError("error")

    trainer = DataParallelTrainer(
        _train_fn,
        scaling_config=ScalingConfig(num_workers=num_workers),
        run_config=RunConfig(
            storage_path=str(tmp_path),
            callbacks=[MyUserCallback()],
        ),
    )
    # The error should NOT be an assertion error from the user callback.
    with pytest.raises(WorkerGroupError):
        trainer.fit()


def run_process_for_sigint_abort(abort_terminates):
    # Lives outside test_sigint_abort because cannot pickle nested functions.

    # Needed to reuse current ray cluster.
    ray.init(address="auto")

    if not abort_terminates:

        async def fake_abort():
            while True:
                pass

        from ray.train.v2._internal.execution.controller import TrainController

        TrainController.abort = fake_abort

    def train_fn():
        signal_actor = ray.get_actor("signal_actor", namespace="test_sigint_abort")
        ray.get(signal_actor.send.remote())
        while True:
            pass

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
    )
    trainer.fit()


@pytest.mark.parametrize(
    "spam_sigint",
    [
        False,
        # Disabling this test because it's flaky.
        # True,
    ],
)
def test_sigint_abort(spam_sigint):
    # Use SignalActor to wait for training to start before sending SIGINT.
    SignalActor = create_remote_signal_actor(ray)
    signal_actor = SignalActor.options(
        name="signal_actor", namespace="test_sigint_abort"
    ).remote()

    # Use spawn because of
    # https://docs.ray.io/en/latest/ray-core/patterns/fork-new-processes.html
    multiprocessing.set_start_method("spawn", force=True)
    process = multiprocessing.Process(
        target=run_process_for_sigint_abort, args=(not spam_sigint,)
    )
    process.start()

    # Wait for training to start.
    ray.get(signal_actor.wait.remote())

    # Verify that process exits after sufficient number of SIGINTS.
    os.kill(process.pid, signal.SIGINT)
    if spam_sigint:
        import time

        assert process.exitcode is None
        # This is flaky. Sometimes SIGINTs are ignored and you need to wait.
        while process.exitcode is None:
            time.sleep(1)
            os.kill(process.pid, signal.SIGINT)
    process.join()


SUPPORTED_METRICS = [
    {"loss": 1.0},
    {"loss": 1, "accuracy": 0.95},
    {"loss": None},
    {"loss": "label"},
    {"nested": {"a": 1}},
]
UNSUPPORTED_METRICS = ["torch_tensor", "nested_torch_tensor", "torch_state_dict"]


def test_supported_report_metrics(tmp_path):
    def train_fn():
        for metric in SUPPORTED_METRICS:
            with tempfile.TemporaryDirectory() as temp_dir:
                ray.train.report(
                    metrics=metric,
                    checkpoint=ray.train.Checkpoint.from_directory(temp_dir),
                )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="test-supported-report-metrics", storage_path=str(tmp_path)
        ),
    )
    result = trainer.fit()
    for (_, actual_metric), expected_metric in zip(
        result.best_checkpoints, SUPPORTED_METRICS, strict=True
    ):
        assert actual_metric == expected_metric

    restored_result = Result.from_path(tmp_path / "test-supported-report-metrics")
    for (_, actual_metric), expected_metric in zip(
        restored_result.best_checkpoints, SUPPORTED_METRICS, strict=True
    ):
        assert actual_metric == expected_metric


@pytest.mark.parametrize("metric_name", UNSUPPORTED_METRICS)
def test_unsupported_report_metrics(metric_name, tmp_path):
    def train_fn():
        if metric_name == "torch_tensor":
            metric = {"loss": torch.tensor(1.0)}
        elif metric_name == "nested_torch_tensor":
            metric = {"nested": {"a": torch.tensor(1.0)}}
        elif metric_name == "torch_state_dict":
            metric = torch.nn.Linear(1, 1).state_dict()
        else:
            raise ValueError()

        with tempfile.TemporaryDirectory() as temp_dir:
            ray.train.report(
                metrics=metric, checkpoint=ray.train.Checkpoint.from_directory(temp_dir)
            )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="test-unsupported-report-metrics", storage_path=str(tmp_path)
        ),
    )
    with pytest.raises(WorkerGroupError) as exc_info:
        trainer.fit()

    assert len(exc_info.value.worker_failures) == 1
    worker_error = exc_info.value.worker_failures[0]
    assert isinstance(worker_error, ValueError)
    assert worker_error.args[0].startswith(
        "Passing objects containing Torch tensors as metrics is not "
        "supported as it will throw an exception on deserialization."
    )


@pytest.mark.parametrize("metric", SUPPORTED_METRICS)
def test_supported_returned_metrics(metric, tmp_path):
    def train_fn():
        return metric

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="test-supported-return-metrics", storage_path=str(tmp_path)
        ),
    )
    result = trainer.fit()
    assert result.return_value == metric


@pytest.mark.parametrize("metric_name", UNSUPPORTED_METRICS)
def test_unsupported_returned_metrics(metric_name, tmp_path):
    def train_fn():
        if metric_name == "torch_tensor":
            metric = {"loss": torch.tensor(1.0)}
        elif metric_name == "nested_torch_tensor":
            metric = {"nested": {"a": torch.tensor(1.0)}}
        elif metric_name == "torch_state_dict":
            metric = torch.nn.Linear(1, 1).state_dict()
        else:
            raise ValueError()

        return metric

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="test-unsupported-report-metrics", storage_path=str(tmp_path)
        ),
    )
    with pytest.raises(WorkerGroupError) as exc_info:
        trainer.fit()

    assert len(exc_info.value.worker_failures) == 1
    worker_error = exc_info.value.worker_failures[0]
    assert isinstance(worker_error, ValueError)
    assert worker_error.args[0].startswith(
        "Returning objects containing Torch tensors from the "
        "training function is not supported as it will throw an "
        "exception on deserialization."
    )


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-x", __file__]))