import os
import sys
import time
from pathlib import Path
from unittest.mock import MagicMock

import pytest

import ray
from ray import tune
from ray.cluster_utils import Cluster
from ray.train._internal.storage import StorageContext
from ray.tune import CheckpointConfig, register_trainable
from ray.tune.error import TuneError
from ray.tune.execution.tune_controller import TuneController
from ray.tune.experiment import Trial
from ray.tune.search import BasicVariantGenerator
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, MyTrainableClass


def _check_trial_running(trial):
    return Path(trial.storage.trial_working_directory, "marker").exists()


def _get_running_trials(runner):
    return [t for t in runner.get_live_trials() if t.status == Trial.RUNNING]


class SlowTrainable(MyTrainableClass):
    def setup(self, config):
        super().setup(config)
        running_marker = Path(self._storage.trial_working_directory, "marker")
        running_marker.touch()

        self._sleep_time = config.get("sleep", 0)

    def step(self):
        time.sleep(self._sleep_time)
        return super().step()


def _start_new_cluster():
    cluster = Cluster(
        initialize_head=True,
        connect=True,
        head_node_args={
            "num_cpus": 1,
            "_system_config": {
                "health_check_initial_delay_ms": 0,
                "health_check_period_ms": 1000,
                "health_check_failure_threshold": 10,
            },
        },
    )
    register_trainable(MOCK_TRAINABLE_NAME, SlowTrainable)
    return cluster


@pytest.fixture
def start_connected_cluster():
    # Start the Ray processes.
    cluster = _start_new_cluster()
    os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
    yield cluster
    # The code after the yield will run as teardown code.
    ray.shutdown()
    cluster.shutdown()


@pytest.fixture
def start_connected_emptyhead_cluster():
    """Starts head with no resources."""
    cluster = Cluster(
        initialize_head=True,
        connect=True,
        head_node_args={
            "num_cpus": 0,
            "_system_config": {
                "health_check_initial_delay_ms": 0,
                "health_check_period_ms": 1000,
                "health_check_failure_threshold": 10,
            },
        },
    )
    os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
    yield cluster
    # The code after the yield will run as teardown code.
    ray.shutdown()
    cluster.shutdown()


@pytest.fixture
def storage(tmp_path):
    os.makedirs(tmp_path / "exp_name" / "trial_name", exist_ok=True)
    yield StorageContext(
        storage_path=str(tmp_path),
        experiment_dir_name="exp_name",
        trial_dir_name="trial_name",
    )


@pytest.fixture(autouse=True)
def register_mock_trainable():
    register_trainable(MOCK_TRAINABLE_NAME, SlowTrainable)
    yield


def test_counting_resources(start_connected_cluster, storage):
    """Tests that Tune accounting is consistent with actual cluster."""

    cluster = start_connected_cluster
    nodes = []
    assert ray.cluster_resources()["CPU"] == 1
    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
    kwargs = {
        "stopping_criterion": {"training_iteration": 10},
        "storage": storage,
        "config": {"sleep": 1},
    }

    trials = [
        Trial(MOCK_TRAINABLE_NAME, **kwargs),
        Trial(MOCK_TRAINABLE_NAME, **kwargs),
    ]
    for t in trials:
        runner.add_trial(t)

    while not any(t.status == Trial.RUNNING for t in trials):
        runner.step()

    running_trials = _get_running_trials(runner)
    assert len(running_trials) == 1
    assert _check_trial_running(running_trials[0])
    assert ray.available_resources().get("CPU", 0) == 0
    nodes += [cluster.add_node(num_cpus=1)]
    cluster.wait_for_nodes()
    assert ray.cluster_resources()["CPU"] == 2
    cluster.remove_node(nodes.pop())
    cluster.wait_for_nodes()
    assert ray.cluster_resources()["CPU"] == 1

    while not any(t.status == Trial.RUNNING for t in trials):
        runner.step()

    # Only 1 trial can be running due to resource limitation.
    assert sum(t.status == Trial.RUNNING for t in runner.get_trials()) == 1

    for i in range(5):
        nodes += [cluster.add_node(num_cpus=1)]
    cluster.wait_for_nodes()
    assert ray.cluster_resources()["CPU"] == 6

    while any(t.status == Trial.PENDING for t in trials):
        runner.step()
    assert sum(t.status == Trial.RUNNING for t in runner.get_trials()) == 2, [
        t.status for t in trials
    ]


def test_trial_processed_after_node_failure(start_connected_emptyhead_cluster, storage):
    """Tests that Tune processes a trial as failed if its node died."""
    cluster = start_connected_emptyhead_cluster
    node = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
    mock_process_failure = MagicMock(side_effect=runner._process_trial_failure)
    runner._process_trial_failure = mock_process_failure
    # Disable recursion in magic mock when saving experiment state
    runner.save_to_dir = lambda *args, **kwargs: None

    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, storage=storage))
    trial = runner.get_trials()[0]

    while trial.status != Trial.RUNNING:
        runner.step()

    assert not mock_process_failure.called

    cluster.remove_node(node)
    while not mock_process_failure.called:
        runner.step()
    assert mock_process_failure.called


def test_remove_node_before_result(start_connected_emptyhead_cluster, storage):
    """Tune continues when node is removed before trial returns."""
    cluster = start_connected_emptyhead_cluster
    node = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
    kwargs = {
        "stopping_criterion": {"training_iteration": 3},
        "checkpoint_config": CheckpointConfig(checkpoint_frequency=2),
        "max_failures": 2,
        "storage": storage,
    }
    trial = Trial(MOCK_TRAINABLE_NAME, **kwargs)
    runner.add_trial(trial)

    while trial.status != Trial.RUNNING:
        runner.step()
    running_trials = _get_running_trials(runner)
    assert len(running_trials) == 1
    assert _check_trial_running(running_trials[0])
    assert not trial.has_reported_at_least_once
    assert trial.status == Trial.RUNNING
    cluster.remove_node(node)
    cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()
    assert ray.cluster_resources()["CPU"] == 1

    while not trial.last_result.get("training_iteration") == 1:
        runner.step()
    assert trial.last_result.get("training_iteration") == 1

    # Process result: discover failure, recover, _train (from scratch)
    while trial.status != Trial.TERMINATED:
        runner.step()

    assert trial.last_result.get("training_iteration") > 1

    with pytest.raises(TuneError):
        runner.step()


def test_trial_requeue(start_connected_emptyhead_cluster, tmpdir, storage):
    """Removing a node in full cluster causes Trial to be requeued."""
    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"

    cluster = start_connected_emptyhead_cluster
    node = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
    kwargs = {
        "stopping_criterion": {"training_iteration": 5},
        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
        "max_failures": 1,
        "storage": storage,
    }

    trials = [
        Trial(MOCK_TRAINABLE_NAME, **kwargs),
        Trial(MOCK_TRAINABLE_NAME, **kwargs),
    ]
    for t in trials:
        runner.add_trial(t)

    while not any(t.status == Trial.RUNNING for t in trials):
        runner.step()

    runner.step()
    runner.step()

    running_trials = _get_running_trials(runner)
    assert len(running_trials) == 1
    assert _check_trial_running(running_trials[0])
    cluster.remove_node(node)
    cluster.wait_for_nodes()
    time.sleep(0.1)  # Sleep so that next step() refreshes cluster resources
    runner.step()  # Process result, dispatch save
    runner.step()  # Process save (detect error), requeue trial
    assert all(t.status == Trial.PENDING for t in trials)


def test_migration_checkpoint_removal(
    start_connected_emptyhead_cluster, tmpdir, storage
):
    """Test checks that trial restarts if checkpoint is lost w/ node fail."""
    cluster = start_connected_emptyhead_cluster
    node = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    runner = TuneController(search_alg=BasicVariantGenerator(), storage=storage)
    kwargs = {
        "stopping_criterion": {"training_iteration": 4},
        "checkpoint_config": CheckpointConfig(checkpoint_frequency=2),
        "max_failures": 2,
        "storage": storage,
    }

    # Test recovery of trial that has been checkpointed
    t1 = Trial(MOCK_TRAINABLE_NAME, **kwargs)
    runner.add_trial(t1)

    # Start trial, process result (x2), process save
    while not t1.has_checkpoint():
        runner.step()

    cluster.add_node(num_cpus=1)
    cluster.remove_node(node)
    cluster.wait_for_nodes()

    while not runner.is_finished():
        runner.step()
    assert t1.status == Trial.TERMINATED


def test_cluster_down_full(start_connected_cluster, tmpdir):
    """Tests that run_experiment restoring works on cluster shutdown."""
    cluster = start_connected_cluster

    base_dict = dict(run=MOCK_TRAINABLE_NAME, stop=dict(training_iteration=3))

    exp1_args = base_dict
    exp2_args = dict(
        base_dict.items(),
        checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
    )
    exp3_args = dict(base_dict.items(), config=dict(mock_error=True))
    exp4_args = dict(
        base_dict.items(),
        config=dict(mock_error=True),
        checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
    )

    all_experiments = {
        "exp1": exp1_args,
        "exp2": exp2_args,
        "exp3": exp3_args,
        "exp4": exp4_args,
    }

    tune.run_experiments(all_experiments, raise_on_failed_trial=False)

    ray.shutdown()
    cluster.shutdown()
    cluster = _start_new_cluster()

    trials = tune.run_experiments(
        all_experiments,
        resume=True,
        raise_on_failed_trial=False,
    )

    assert len(trials) == 4
    assert all(t.status in [Trial.TERMINATED, Trial.ERROR] for t in trials)
    ray.shutdown()
    cluster.shutdown()


if __name__ == "__main__":
    import pytest

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