import json
import logging
import os
import sys
import tempfile
import time
from unittest import mock

import pytest

import ray
from ray.air.constants import TRAINING_ITERATION
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
from ray.train._internal.session import _TrainingResult
from ray.train._internal.storage import StorageContext
from ray.train.tests.util import mock_storage_context
from ray.tune import (
    Callback,
    Checkpoint,
    CheckpointConfig,
    PlacementGroupFactory,
    ResumeConfig,
)
from ray.tune.execution.tune_controller import TuneController
from ray.tune.experiment import Trial
from ray.tune.result import DONE
from ray.tune.schedulers import FIFOScheduler
from ray.tune.search import BasicVariantGenerator
from ray.tune.tests.tune_test_util import TrialResultObserver
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable

STORAGE = mock_storage_context()


@pytest.fixture(autouse=True)
def register_test_trainable():
    register_mock_trainable()


@pytest.fixture(scope="function")
def ray_start_4_cpus_2_gpus_extra():
    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
    yield address_info
    ray.shutdown()


def create_mock_components():
    class _MockScheduler(FIFOScheduler):
        errored_trials = []

        def on_trial_error(self, tune_controller, trial):
            self.errored_trials += [trial]

    class _MockSearchAlg(BasicVariantGenerator):
        errored_trials = []

        def on_trial_complete(self, trial_id, error=False, **kwargs):
            if error:
                self.errored_trials += [trial_id]

    searchalg = _MockSearchAlg()
    scheduler = _MockScheduler()
    return searchalg, scheduler


def num_checkpoints(trial):
    return sum(
        item.startswith("checkpoint_")
        for item in os.listdir(trial.storage.trial_fs_path)
    )


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_save_restore(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
):
    """Test that a checkpoint is saved and can be used to restore a trainable.

    The trainable saves a checkpoint and terminates. We then start another trial
    that should restore from the saved checkpoint and assert that it picks up
    the state and continues to run to termination.

    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testCheckpointing
    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testRestoreMetricsAfterCheckpointing  # noqa
    """
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 1},
        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
        "storage": STORAGE,
    }
    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
    trials = runner.get_trials()

    runner.step()  # Start trial

    while trials[0].status != Trial.RUNNING:
        runner.step()

    while trials[0].status != Trial.TERMINATED:
        runner.step()

    assert trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION] == 1
    assert trials[0].last_result[TRAINING_ITERATION] == 1
    assert trials[0].last_result["iterations_since_restore"] == 1

    # Prepare new trial
    kwargs["restore_path"] = trials[0].checkpoint.path
    new_trial = Trial(MOCK_TRAINABLE_NAME, **kwargs)
    runner.add_trial(new_trial)
    trials = runner.get_trials()

    assert trials[1].status == Trial.PENDING

    # Start trial, restore, run to termination
    while trials[1].status != Trial.RUNNING:
        runner.step()

    # Restore
    runner.step()

    # Run to termination
    while trials[1].status != Trial.TERMINATED:
        runner.step()

    assert trials[0].latest_checkpoint_result.metrics[TRAINING_ITERATION] == 1
    assert trials[1].last_result[TRAINING_ITERATION] == 1
    assert trials[1].last_result["iterations_since_restore"] == 1


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_at_end(ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir):
    """Test that a checkpoint is saved at end for class trainables with that config.

    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testCheckpointingAtEnd
    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testResultDone
    """
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        storage=STORAGE,
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 2},
        "checkpoint_config": CheckpointConfig(checkpoint_at_end=True),
        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
        "storage": STORAGE,
    }
    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
    trials = runner.get_trials()

    while not runner.is_finished():
        runner.step()

    assert trials[0].has_checkpoint()
    assert trials[0].last_result[DONE]


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_pause_resume_trial(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
):
    """Test that trial that is paused and resumed picks up its last checkpoint.

    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testPauseThenResume
    """
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        storage=STORAGE,
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 2},
        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
        "checkpoint_config": CheckpointConfig(checkpoint_frequency=1),
        "storage": STORAGE,
    }
    runner.add_trial(Trial(MOCK_TRAINABLE_NAME, **kwargs))
    trials = runner.get_trials()

    while trials[0].status != Trial.RUNNING:
        runner.step()

    runner._schedule_trial_pause(trials[0], should_checkpoint=True)

    while trials[0].status != Trial.PAUSED:
        runner.step()

    assert trials[0].has_checkpoint()
    assert not trials[0].last_result.get(DONE), trials[0].last_result

    # Start again
    runner._set_trial_status(trials[0], Trial.PENDING)

    while trials[0].status != Trial.RUNNING:
        runner.step()

    while trials[0].status != Trial.TERMINATED:
        runner.step()

    assert trials[0].checkpoint
    assert trials[0].last_result[TRAINING_ITERATION] == 2
    assert trials[0].last_result["iterations_since_restore"] == 1
    assert trials[0].last_result["time_since_restore"] > 0


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_num_to_keep(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
):
    """Test that only num_to_keep checkpoints are kept.

    This should also hold true when the experiment is resumed.

    Legacy test: test_trial_runner_2.py::TrialRunnerTest::testPauseResumeCheckpointCount
    """
    trial = Trial(
        MOCK_TRAINABLE_NAME,
        checkpoint_config=CheckpointConfig(num_to_keep=2),
        storage=STORAGE,
    )
    trial.init_local_path()

    def write_checkpoint(trial: Trial, index: int):
        checkpoint_dir = tmp_path / StorageContext._make_checkpoint_dir_name(index)
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
        result = {"training_iteration": index}
        with open(os.path.join(checkpoint_dir, "cp.json"), "w") as f:
            json.dump(result, f)

        checkpoint = Checkpoint.from_directory(checkpoint_dir)
        return _TrainingResult(checkpoint=checkpoint, metrics=result)

    def get_checkpoint_dirs(trial: Trial):
        return [d for d in os.listdir(tmp_path) if d.startswith("checkpoint_")]

    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
    )

    runner.add_trial(trial)

    # Write 1 checkpoint
    result = write_checkpoint(trial, 1)
    runner._on_saving_result(trial, result)

    # Expect 1 checkpoint
    cp_dirs = get_checkpoint_dirs(trial)
    assert len(cp_dirs) == 1, f"Checkpoint dirs: {cp_dirs}"

    # Write second checkpoint
    result = write_checkpoint(trial, 2)
    runner._on_saving_result(trial, result)

    # Expect 2 checkpoints
    cp_dirs = get_checkpoint_dirs(trial)
    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"

    # Write third checkpoint
    result = write_checkpoint(trial, 3)
    runner._on_saving_result(trial, result)

    # Expect 2 checkpoints because num_to_keep = 2
    cp_dirs = get_checkpoint_dirs(trial)
    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"

    # Re-instantiate trial runner and resume
    runner.checkpoint(force=True, wait=True)
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        storage=STORAGE,
        resume_config=ResumeConfig(),
    )

    trial = runner.get_trials()[0]

    # Write fourth checkpoint
    result = write_checkpoint(trial, 4)
    runner._on_saving_result(trial, result)

    # Expect 2 checkpoints because num_to_keep = 2
    cp_dirs = get_checkpoint_dirs(trial)
    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"

    # Write fifth checkpoint
    result = write_checkpoint(trial, 5)
    runner._on_saving_result(trial, result)

    # Expect 2 checkpoints because num_to_keep = 2
    cp_dirs = get_checkpoint_dirs(trial)
    assert len(cp_dirs) == 2, f"Checkpoint dirs: {cp_dirs}"

    # Checkpoints before restore should be deleted
    assert "checkpoint_000004" in cp_dirs
    assert "checkpoint_000005" in cp_dirs

    assert "checkpoint_000002" not in cp_dirs
    assert "checkpoint_000003" not in cp_dirs


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_freq_buffered(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
):
    """Test that trial checkpoints are a lower bound for buffered training iterations.

    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointFreqBuffered
    """
    with mock.patch.dict(
        os.environ,
        {"TUNE_RESULT_BUFFER_LENGTH": "7", "TUNE_RESULT_BUFFER_MIN_TIME_S": "1"},
    ):
        trial = Trial(
            MOCK_TRAINABLE_NAME,
            checkpoint_config=CheckpointConfig(checkpoint_frequency=3),
            storage=STORAGE,
        )
        runner = TuneController(
            resource_manager_factory=lambda: resource_manager_cls(),
            storage=STORAGE,
            checkpoint_period=0,
        )
        runner.add_trial(trial)

        while not trial.is_saving:
            runner.step()
        runner.step()
        assert trial.last_result[TRAINING_ITERATION] == 3
        assert num_checkpoints(trial) == 1

        while not trial.is_saving:
            runner.step()
        runner.step()
        assert trial.last_result[TRAINING_ITERATION] == 6
        assert num_checkpoints(trial) == 2

        while not trial.is_saving:
            runner.step()
        runner.step()
        assert trial.last_result[TRAINING_ITERATION] == 9
        assert num_checkpoints(trial) == 3


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_at_end_not_buffered(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
):
    """Test that trials with `checkpoint_at_end=True` are never buffered.

    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointAtEndNotBuffered
    """
    with mock.patch.dict(
        os.environ,
        {"TUNE_RESULT_BUFFER_LENGTH": "7", "TUNE_RESULT_BUFFER_MIN_TIME_S": "0.5"},
    ):
        trial = Trial(
            MOCK_TRAINABLE_NAME,
            checkpoint_config=CheckpointConfig(
                checkpoint_at_end=True,
            ),
            stopping_criterion={"training_iteration": 4},
            storage=STORAGE,
        )
        observer = TrialResultObserver()
        runner = TuneController(
            resource_manager_factory=lambda: resource_manager_cls(),
            storage=STORAGE,
            callbacks=[observer],
        )
        runner.add_trial(trial)

        while not observer.just_received_a_result():
            runner.step()
        assert trial.last_result[TRAINING_ITERATION] == 1
        assert num_checkpoints(trial) == 0

        while True:
            runner.step()
            if observer.just_received_a_result():
                break
        assert trial.last_result[TRAINING_ITERATION] == 2
        assert num_checkpoints(trial) == 0

        while True:
            runner.step()
            if observer.just_received_a_result():
                break
        assert trial.last_result[TRAINING_ITERATION] == 3
        assert num_checkpoints(trial) == 0

        while True:
            runner.step()
            if observer.just_received_a_result():
                break
        assert trial.last_result[TRAINING_ITERATION] == 4

        while not runner.is_finished():
            runner.step()
        assert num_checkpoints(trial) == 1


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_checkpoint_auto_period(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
):
    """Test that the checkpoint auto period is adjusted when syncing takes a long time.

    Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointAutoPeriod
    """
    storage = mock_storage_context()

    with tempfile.TemporaryDirectory() as local_dir:
        storage.storage_local_path = local_dir

        runner = TuneController(
            resource_manager_factory=lambda: resource_manager_cls(),
            storage=storage,
            checkpoint_period="auto",
        )

        with mock.patch.object(runner, "save_to_dir") as save_to_dir:
            save_to_dir.side_effect = lambda *a, **kw: time.sleep(2)
            runner.add_trial(Trial(MOCK_TRAINABLE_NAME, storage=storage))
            runner.step()  # Run one step, this will trigger checkpointing

        assert runner._checkpoint_manager._checkpoint_period > 38.0


def test_checkpoint_force_with_num_to_keep(ray_start_4_cpus_2_gpus_extra, tmp_path):
    """Test that cloud syncing is forced if one of the trials has made more
    than num_to_keep checkpoints since last sync.
    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
        testCloudCheckpointForceWithNumToKeep
    """
    storage = mock_storage_context()
    # Needed to avoid infinite recursion error on CI runners
    storage.syncer.__getstate__ = lambda *a, **kw: {}

    with mock.patch.object(storage.syncer, "sync_up") as sync_up:
        num_to_keep = 2
        checkpoint_config = CheckpointConfig(
            num_to_keep=num_to_keep, checkpoint_frequency=1
        )

        runner = TuneController(
            resource_manager_factory=lambda: PlacementGroupResourceManager(),
            storage=storage,
            checkpoint_period=100,  # only rely on force syncing
            trial_checkpoint_config=checkpoint_config,
        )

        class CheckpointingTrial(Trial):
            def should_checkpoint(self):
                return True

            def get_json_state(self):
                return "", ""

        trial = CheckpointingTrial(
            MOCK_TRAINABLE_NAME,
            checkpoint_config=checkpoint_config,
            stopping_criterion={"training_iteration": 10},
            storage=storage,
        )
        runner.add_trial(trial)

        # also check if the warning is printed
        buffer = []
        from ray.tune.execution.experiment_state import logger

        with mock.patch.object(logger, "warning", lambda x: buffer.append(x)):
            while not runner.is_finished():
                runner.step()

        assert any(
            "Experiment state snapshotting has been triggered multiple times" in x
            for x in buffer
        )
        # We should sync 6 times:
        # The first checkpoint happens when the experiment starts,
        # since no checkpoints have happened yet
        # (This corresponds to the new_trial event in the runner loop)
        # Then, every num_to_keep=2 checkpoints, we should perform a forced checkpoint
        # which results in 5 more checkpoints (running for 10 iterations),
        # giving a total of 6
        assert sync_up.call_count == 6


def test_checkpoint_force_by_trial_callback(ray_start_4_cpus_2_gpus_extra, tmp_path):
    """Test that cloud syncing is forced if one of the trials has made more
    than num_to_keep checkpoints since last sync.
    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
        testCloudCheckpointForceWithNumToKeep
    """

    class CheckpointCallback(Callback):
        def __init__(self):
            self.num_checkpoints = 0

        def on_trial_result(self, iteration, trials, trial: Trial, result, **info):
            # Checkpoint every two iterations
            if result[TRAINING_ITERATION] % 2 == 0:
                self.num_checkpoints += 1
                result["should_checkpoint"] = True

    storage = mock_storage_context()

    # disable automatic checkpointing
    checkpoint_config = CheckpointConfig(checkpoint_frequency=0)
    callback = CheckpointCallback()
    runner = TuneController(
        resource_manager_factory=PlacementGroupResourceManager,
        storage=storage,
        callbacks=[callback],
        trial_checkpoint_config=checkpoint_config,
    )

    trial = Trial(
        MOCK_TRAINABLE_NAME,
        checkpoint_config=checkpoint_config,
        stopping_criterion={"training_iteration": 6},
        storage=storage,
    )
    runner.add_trial(trial)

    while not runner.is_finished():
        runner.step()

    assert callback.num_checkpoints == 3
    assert num_checkpoints(trial) == 3


def test_checkpoint_sync_up_timeout(
    ray_start_4_cpus_2_gpus_extra, tmp_path, monkeypatch
):
    """Test that trial runner experiment checkpointing times out correctly.

    Legacy test: test_trial_runner_3.py::TrialRunnerTest::
        testForcedCloudCheckpointSyncTimeout
    """
    storage = mock_storage_context(sync_config=ray.tune.SyncConfig(sync_timeout=0.5))
    monkeypatch.setenv("TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "0.25")

    def _hanging_upload_to_fs_path(*args, **kwargs):
        time.sleep(200)

    monkeypatch.setattr(
        ray.train._internal.storage,
        "_upload_to_fs_path",
        _hanging_upload_to_fs_path,
    )

    runner = TuneController(
        resource_manager_factory=lambda: PlacementGroupResourceManager(),
        storage=storage,
    )

    # Start a hanging sync that should not block the controller
    runner.checkpoint()

    buffer = []
    logger = logging.getLogger("ray.tune.execution.experiment_state")
    with mock.patch.object(logger, "error", lambda x, **kwargs: buffer.append(x)):
        with mock.patch.object(logger, "warning", lambda x: buffer.append(x)):
            runner.checkpoint(force=True, wait=True)

    # We should see a log about the timeout
    assert any("Saving experiment state to storage" in x for x in buffer)
    # We should also have a warning about the slow upload
    assert any("may be a performance bottleneck" in x for x in buffer)


def test_checkpoint_sync_up_error(ray_start_4_cpus_2_gpus_extra, tmp_path, monkeypatch):
    """Test that trial runner experiment checkpointing handles errors correctly."""
    storage = mock_storage_context()

    def _failing_upload_to_fs_path(*args, **kwargs):
        raise RuntimeError("Upload failing...")

    monkeypatch.setattr(
        ray.train._internal.storage,
        "_upload_to_fs_path",
        _failing_upload_to_fs_path,
    )

    runner = TuneController(
        resource_manager_factory=lambda: PlacementGroupResourceManager(),
        storage=storage,
    )

    # Launching a failing upload task should not crash the controller / main thread
    runner.checkpoint()

    buffer = []
    logger = logging.getLogger("ray.tune.execution.experiment_state")
    with mock.patch.object(logger, "error", lambda x, **kwargs: buffer.append(x)):
        runner.checkpoint(force=True)

    # We should see a log about the failure
    assert any("Saving experiment state to storage" in x for x in buffer)


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))